diff --git a/CHANGELOG.md b/CHANGELOG.md index a8640fd..b0ebb7a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,25 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Added + +- Sibling requests that arrive together no longer each prefill the shared + prefix cold: the server admits the first one, waits for its stores, and + starts the rest warm (`GMLX_APC_FRESH_WAIT_MS`, `0` disables). +- Requests beyond the queue cap now get an immediate 503 with Retry-After + instead of queueing toward the timeout (`GMLX_QUEUE_DEPTH_CAP`). +- A prompt that cannot fit in memory now gets a 400 with the numbers before + the stream opens, instead of dying mid-stream (`GMLX_PREFLIGHT_MEM=0` + disables). +- `seed` is now honored per request inside a batch. Before, only the first + request's seed took effect and it colored every row. +- Short prompts on sliding-window models now cache their block-aligned + prefix at retirement instead of nothing, so an immediate follow-up + turn starts warm. +- Decode concurrency is now a control (`GMLX_DECODE_BATCH`, default 8). + The server always decoded up to 32 requests together, which slows every + stream past the width where total throughput stops growing. + ## [0.3.2] - 2026-08-13 ### Changed diff --git a/docs/server-config.md b/docs/server-config.md index 5822c2f..6cd0cfc 100644 --- a/docs/server-config.md +++ b/docs/server-config.md @@ -309,6 +309,28 @@ deep-context safety case. See policy, the `GMLX_CACHE_LIMIT_GB` env override (env wins over this key), and the explicit-unlimited escape. +A text request whose prompt alone cannot fit in memory gets an immediate +HTTP 400 with the estimated need and the available budget in the body, +instead of dying mid-stream. The estimate prices prompt KV at the model's +per-token cost (GQA heads, MLA latents, sliding windows, and quantized KV +all lower it) plus the prefill score transient, against the working set +with the batch drained. `max_tokens` counts only when the request pins it +explicitly; default-max requests are never rejected on generation length. +Media requests are not estimated in v1. `GMLX_PREFLIGHT_MEM=0` disables. + +Decode concurrency (how many requests generate tokens together in one batch +step) defaults to 8; past that width aggregate throughput gains shrink while +every stream slows. `GMLX_DECODE_BATCH` sets it (`0` restores the upstream +default of 32). + +Requests beyond the waiting-queue cap get an immediate HTTP 503 with a +`Retry-After` header instead of queueing toward the token-queue timeout. The +JSON body names the cap and the current depth; the header value is the +estimated drain time, clamped to 2-60 seconds. Harness SDKs back off on 503 +and retry, which beats holding a silent socket for half an hour. +`GMLX_QUEUE_DEPTH_CAP` sets the cap (default 2 x the decode concurrency; +`0` disables the check). + While a streaming request is silent (most notably during that long prefill), the server emits an SSE comment line (`: keepalive`) every 15 seconds so clients with a between-bytes read timeout don't drop the connection before @@ -917,6 +939,18 @@ The request fields mlx-vlm already honours, carried verbatim into generation: `repetition_context_size`, `enable_thinking`, `thinking_budget`, `thinking_start_token`, `thinking_end_token`. +`seed` is honored per request, in-batch: each seeded request draws its +tokens from its own key stream while unseeded rows in the same batch keep +the shared stream, byte for byte. Seed guarantees a deterministic sampling +stream for that request. It does not guarantee bitwise-identical output +across runs with different batch composition, because batched matmul +reduction order shifts logits at float tolerance; the same composition +(for example a solo replay) reproduces exactly. With speculative decoding, +drafts are greedy and a seeded single-stream request's target draws come +from the same per-request key stream, so a replay matches only across runs +with the same speculation setting; seeded rows inside a batched +speculative decode fall back to the shared stream. + `thinking_start_token` / `thinking_end_token` override the `` / `` defaults everywhere the server needs the model's real reasoning markers (open-think detection, `thinking_budget`, the streamed @@ -1183,6 +1217,8 @@ and store counts surface on the authed `GET /v1/metrics`. | `GMLX_APC_CKPT_BUDGET_MB` | Byte budget for checkpoint-record payload (recurrent states + KV tails), in MB (default `4096`). A GDN record can carry >100 MB of state and each request saves several checkpoints, so expect resident memory to grow toward this budget on hybrid models under sustained multi-turn traffic; lower it if 4 GB of cache is too much for your machine. | | `GMLX_APC_DECODE_CKPT` | Decode-time snapshot interval in generated tokens on hybrid models, anchored to the prompt end (default `512`; `0` off; widens automatically with context). | | `GMLX_APC_RETIRE_LCP` | `0` keys retirement on the forwarded ids instead of the predicted next-turn render (also disables decode-time snapshots, which key on the prediction). | +| `GMLX_APC_FRESH_WAIT_MS` | Hold ceiling for the freshness admission gate, in ms (default `500`; `0` disables the gate). Sibling requests that arrive together admit one formation apart instead of together and cold: the first request prefills and stores the shared prefix, and the held siblings then admit warm. A sibling held past the ceiling admits cold. | +| `GMLX_APC_FRESH_MIN` | Minimum uncovered shared-prefix tokens before the gate holds a sibling (default `256`). Below the floor the duplicate prefill costs less than the wait. | | `GMLX_FAITHFUL_HISTORY` | `0` restores mlx-vlm's stock chat-history rebuild, which drops `reasoning_content` from non-tool assistant messages before the template sees it (see `chat_template_kwargs`). | --- diff --git a/gmlx/apc_manager.py b/gmlx/apc_manager.py index 510caaa..6713471 100644 --- a/gmlx/apc_manager.py +++ b/gmlx/apc_manager.py @@ -86,14 +86,18 @@ def stats_snapshot(self) -> dict: wrap: super() + merge). Visible at /v1/cache/stats -- a ckpt model with zeroed ckpt_* keys is broken, not idle.""" from .cache_snapshot import ckpt_stats_snapshot + from .prefix_cache import spec_prefix_stats snap = super().stats_snapshot() snap.update(ckpt_stats_snapshot(self)) + snap.update(spec_prefix_stats()) return snap def reset_stats(self) -> None: from .cache_snapshot import ckpt_stats_clear + from .prefix_cache import spec_prefix_stats_clear super().reset_stats() ckpt_stats_clear(self) + spec_prefix_stats_clear() def clear(self) -> None: """Stock clear plus the ckpt tier: the pool wipe zeroes the block @@ -102,9 +106,13 @@ def clear(self) -> None: lookup can never pin a record between the pool wipe and the record drop.""" from .cache_snapshot import ckpt_reset + from .prefix_cache import clear_all_spec_prefix_caches with self.lock: super().clear() ckpt_reset(self) + # Outside the manager lock: the spec prefix cache has its own + # lifetime and its pinned snapshots must not survive a reset. + clear_all_spec_prefix_caches() def store_ckpt_blocks(self, token_ids, layer_keys, layer_values, *, extra_hash=0, disk=True): diff --git a/gmlx/cache_snapshot.py b/gmlx/cache_snapshot.py index 8a8c30d..f5a5e33 100644 --- a/gmlx/cache_snapshot.py +++ b/gmlx/cache_snapshot.py @@ -756,7 +756,7 @@ def _ckpt_records(manager) -> "OrderedDict": "ckpt_stores", "ckpt_hits", "ckpt_matched_tokens", "ckpt_missed_adoptions", "ckpt_skeleton_writes", "sidecar_writes", "retire_fallback_full", "ckpt_pool_evictions", - "anchor_stores", "anchor_hits", + "ckpt_grid_truncate", "anchor_stores", "anchor_hits", ) @@ -1034,7 +1034,8 @@ def ckpt_store( extra_hash: int = 0, skeleton_disk: bool = True, kind: str = "boundary", -) -> bool: + grid_truncate: bool = False, +) -> int: """Store a hybrid checkpoint at ``p = len(token_ids)``. Single-row cache list, KV/rotating offsets == p. Plain KV rides the @@ -1045,10 +1046,16 @@ def ckpt_store( skeleton inlines recurrent state, >100 MB per GDN checkpoint -- interval boundaries superseded minutes later do not earn that). ``kind`` stamps the record's retention class (see _CkptRecord). - Never raises. + ``grid_truncate`` turns the below-window off-grid rotating decline + into a terminal store at the largest block-aligned prefix: pre-wrap + the buffer is a temporal prefix, so a slice is a faithful shorter + run. Non-recurrent layouts only (state cannot rewind), memory-only + (the live cache's offset would stamp a mismatched skeleton). + Returns the stored length in tokens, 0 when nothing stored. Never + raises. """ if manager is None or token_ids is None: - return False + return 0 from .cache_compat import cache_types, runtime_cache_module kv_types = cache_types("KVCache") @@ -1062,7 +1069,7 @@ def ckpt_store( layout = ckpt_layout(prompt_cache, bs) if p < 2 or layout is None: _ckpt_decline(manager, "layout") - return False + return 0 if kind == "replay" and "arr" in layout: # The disk path knows no kinds, so a skeleton here would let # a restart serve this record past the replay adopt gate -- @@ -1081,7 +1088,7 @@ def ckpt_store( "APC ckpt store declined: BufferedRotatingKVCache rows " "cannot snapshot (support deferred)") _ckpt_decline(manager, "buffered") - return False + return 0 for c in prompt_cache: off = getattr(c, "offset", None) if off is not None and not isinstance(c, rot_types) \ @@ -1089,14 +1096,15 @@ def ckpt_store( _log.info("APC ckpt store skipped: KV offset %d != %d", int(off), p) _ckpt_decline(manager, "offset") - return False + return 0 if isinstance(c, rot_types) and int(c.offset) != p: _log.info("APC ckpt store skipped: rot offset %d != %d", int(c.offset), p) _ckpt_decline(manager, "offset") - return False + return 0 has_rot = any(_is_rot(t) for t in layout) b_full = _ckpt_block_prefix(p, bs) + trunc_from = None if has_rot and b_full != p: # Off-grid p is storable once the window has wrapped: the # canonical window is then exactly W tokens -- whole blocks @@ -1111,14 +1119,32 @@ def ckpt_store( _log.warning("APC ckpt store declined: rotating geometry " "unavailable at grid gate") _ckpt_decline(manager, "layout") - return False + return 0 if p < geom[0]: + has_arr = any(not isinstance(c, (kv_types, rot_types)) + for c in prompt_cache) + if not grid_truncate or has_arr or b_full < 2: + _log.info( + "APC ckpt store declined: off-grid rotating store " + "below the window (p=%d < W=%d, %d %% %d != 0)", + p, geom[0], p, bs) + _ckpt_decline(manager, "grid") + return 0 + # Terminal grid store: pre-wrap the buffer is a temporal + # prefix, so the block-aligned slice is a faithful + # shorter run. Memory-only: a skeleton would stamp the + # live cache's deeper offset. + trunc_from = p + p = b_full + ids = ids[:p] + # No skeleton and no window-chain disk blocks: without + # the skeleton nothing re-indexes them after a restart. + skeleton_disk = False + rot_disk = False + _ckpt_bump(manager, "ckpt_grid_truncate") _log.info( - "APC ckpt store declined: off-grid rotating store " - "below the window (p=%d < W=%d, %d %% %d != 0)", - p, geom[0], p, bs) - _ckpt_decline(manager, "grid") - return False + "APC ckpt store: terminal grid store at %d (prompt " + "%d below window %d)", p, trunc_from, geom[0]) tail_len = p - b_full salted = ckpt_extra_hash(extra_hash) kv_caches = [c for c in prompt_cache if isinstance(c, kv_types) @@ -1185,21 +1211,29 @@ def _once(): ids[:b_full], lk, lv, extra=salted, disk=True, need=b_full // bs, what="main") if got_main is None: - return False + return 0 main_blocks = got_main rot_meta = None if rot_caches: canon = [rotating_canonical_window(c) for c in rot_caches] + if trunc_from is not None: + # Slice each canonical window to the terminal grid p and + # restamp its meta as the shorter run's canonical form + # (pre-wrap: L == offset == p, idx == L). + canon = [None if cw is None else + (cw[0][..., :p, :], cw[1][..., :p, :], + (cw[2][0], cw[2][1], p, p)) + for cw in canon] if any(cw is None for cw in canon): _ckpt_decline(manager, "canon") manager.release(main_blocks) - return False + return 0 metas = {cw[2] for cw in canon} if len(metas) != 1: _ckpt_decline(manager, "canon") manager.release(main_blocks) - return False + return 0 rot_meta = canon[0][2] keep, _w, _off, L = rot_meta canon_ids = ids[:keep] + ids[p - (L - keep):p] @@ -1214,7 +1248,7 @@ def _once(): need=L // bs, what="window") if got_win is None: manager.release(main_blocks) - return False + return 0 bounded_blocks = got_win states = [_clone_single_row(c) for c in arr_caches] @@ -1222,7 +1256,7 @@ def _once(): _ckpt_decline(manager, "clone") manager.release(main_blocks) manager.release(bounded_blocks) - return False + return 0 tails = None if tail_len and kv_caches: tails = [] @@ -1235,7 +1269,7 @@ def _once(): _ckpt_decline(manager, "clone") manager.release(main_blocks) manager.release(bounded_blocks) - return False + return 0 tails.append(t) rec = _CkptRecord( @@ -1255,7 +1289,7 @@ def _once(): "APC ckpt store: tokens=%d main=%d window=%d tail=%d states=%d", p, len(rec.main_blocks), len(rec.bounded_blocks), tail_len, len(states)) - return True + return p except Exception: try: _ckpt_decline(manager, "exception") @@ -1264,7 +1298,7 @@ def _once(): except Exception: pass # best-effort release on the failure path _log.warning("APC ckpt store failed; continuing", exc_info=True) - return False + return 0 def _ckpt_disk_write(manager, ids, prompt_cache, layout, p, b_full, @@ -1854,9 +1888,9 @@ def _snap_assemble(prompt_cache: list[Any], states: list[Any], def _ckpt_retirement(manager, ids, prompt_cache, *, extra_hash, max_len, decode_snaps) -> int: - """Ckpt-mode retirement: the full sequence when it can store whole, - else the newest decode-time snapshot at or below the replayable - prefix. Never spills to the exact tier -- on ckpt models the exact + """Ckpt-mode retirement: the full sequence when it can store whole + (a short rotating prompt truncates to the block grid), else the + newest decode-time snapshot at or below the replayable prefix. Never spills to the exact tier -- on ckpt models the exact tier stays empty, so the stock warm path never bypasses arming. Returns the stored length (0 = nothing).""" try: @@ -1867,12 +1901,14 @@ def _ckpt_retirement(manager, ids, prompt_cache, *, extra_hash, # The row is already single-row on the B=1 path; ckpt_store # slices it directly (its own stores copy internally), so no # full-cache clone happens -- the exact tier's whole sin. A - # rotating layer declines here below the window (a sub-wrap - # store needs the block grid); the ring below holds aligned - # clones for exactly that case. - if ckpt_store(manager, ids, prompt_cache, - extra_hash=extra_hash, kind="retire"): - return len(ids) + # rotating layer below the window stores its block-grid + # prefix (grid_truncate); the ring below holds aligned + # clones for the post-wrap off-grid cases. + stored = ckpt_store(manager, ids, prompt_cache, + extra_hash=extra_hash, + grid_truncate=True, kind="retire") + if stored: + return stored for p, states in sorted(decode_snaps or (), key=lambda s: s[0], reverse=True): if not 2 <= p <= cap: @@ -1893,13 +1929,15 @@ def _ckpt_retirement(manager, ids, prompt_cache, *, extra_hash, # Reason-counted so fallback traffic is distinguishable from # turn reuse. Only when the whole-sequence branch above did not # already try (cap == len means it ran and declined). - if cap < len(ids) and len(ids) >= 2 and ckpt_store( - manager, ids, prompt_cache, extra_hash=extra_hash, - kind="retire"): - _ckpt_bump(manager, "retire_fallback_full") - _log.info("APC retirement: full-sequence fallback stored at " - "%d (replayable prefix %d)", len(ids), cap) - return len(ids) + if cap < len(ids) and len(ids) >= 2: + stored = ckpt_store(manager, ids, prompt_cache, + extra_hash=extra_hash, + grid_truncate=True, kind="retire") + if stored: + _ckpt_bump(manager, "retire_fallback_full") + _log.info("APC retirement: full-sequence fallback stored " + "at %d (replayable prefix %d)", stored, cap) + return stored _log.info("APC retirement skipped: no decode snapshot at or " "below the replayable prefix %d (full %d)", cap, len(ids)) diff --git a/gmlx/decode_batch.py b/gmlx/decode_batch.py new file mode 100644 index 0000000..c26a52e --- /dev/null +++ b/gmlx/decode_batch.py @@ -0,0 +1,39 @@ +"""Decode concurrency control (GMLX_DECODE_BATCH). + +Bounds how many requests decode together in one batch step. Upstream +hard-wires 32; aggregate throughput saturates well below that at depth +while per-stream latency keeps degrading. + +Knobs: + GMLX_DECODE_BATCH decode concurrency (default 8; 0 = upstream + default) +""" + +from __future__ import annotations + +import os + +DEFAULT_DECODE_BATCH = 8 + + +def _upstream_default() -> int: + try: + from mlx_vlm.generate.ar import DEFAULT_COMPLETION_BATCH_SIZE + return int(DEFAULT_COMPLETION_BATCH_SIZE) + except Exception: + return 32 + + +def decode_batch() -> int: + """Effective serve-path decode concurrency, always positive.""" + raw = os.environ.get("GMLX_DECODE_BATCH", "").strip() + if raw: + try: + v = int(raw) + except ValueError: + return DEFAULT_DECODE_BATCH + if v > 0: + return v + if v == 0: + return _upstream_default() + return DEFAULT_DECODE_BATCH diff --git a/gmlx/deepseek_v4_cache.py b/gmlx/deepseek_v4_cache.py index 8320e6a..deb780b 100644 --- a/gmlx/deepseek_v4_cache.py +++ b/gmlx/deepseek_v4_cache.py @@ -563,7 +563,9 @@ def accumulate_windows(self, kv: mx.array, gate: mx.array, offset): # the stashed objects keep the pre-update contents. The pooled # tensor needs no snapshot: update_and_fetch only writes beyond the # old _pool_lengths, so restoring the length lists is enough. - if L <= 2: + # L bound matches the scalar cache: MTP-verify sized updates only + # (block-total-B rounds verify B-1+1 wide, dspark up to 6). + if L <= 6: self._undo = ( self.buf_kv, self.buf_gate, @@ -774,14 +776,19 @@ def is_trimmable(self): return True return self._can_undo(1) + def _can_trim(self, n): + """n-aware trimmability probe (MTP rollback two-phase check).""" + if self.pooled is None or n <= min(self.remainder): + return True + return self._can_undo(n) + def _can_undo(self, n): + # A replay that re-completes windows is a per-row watermark move + # (the pooled rows the undone update appended are identical and + # sit above the restored watermark), so any n within the stashed + # update is undoable. undo = self._undo - if undo is None: - return False - k = undo[5].shape[1] - n - # The replayed confirmed prefix must stay inside the buffer for - # every row (a replay that pools again cannot be reconstructed). - return k >= 0 and all(r + k < self.ratio for r in undo[2]) + return undo is not None and undo[5].shape[1] - n >= 0 def trim(self, n): if n <= min(self.remainder): @@ -807,11 +814,56 @@ def trim(self, n): self.remainder = list(remainder) self._pool_lengths = list(pool_lengths) self._processed = list(processed) - if k > 0: - # Replay the confirmed prefix; _can_undo guarantees it stays in - # the buffer, so no window is recompressed. + if k <= 0: + return n + if all(r + k < self.ratio for r in self.remainder): + # Replay stays inside every row's buffer; no window completes. self.accumulate_windows(kv[:, :k], gate[:, :k], 0) self._undo = None + return n + # The confirmed prefix re-completes window(s) for at least one + # row. Their pooled rows are exactly the first rows the undone + # update appended (identical inputs and lookback) and still sit + # above the restored watermark, so the replay is a per-row + # watermark move; remainder buffer and lookback rebuild from the + # stashed raw inputs -- no recompression. + ratio = self.ratio + B = kv.shape[0] + for i in range(B): + r = self.remainder[i] + total = r + k + w = total // ratio + new_rem = total % ratio + if r > 0: + seq_kv = mx.concatenate( + [self.buf_kv[i : i + 1, :r], kv[i : i + 1, :k]], axis=1) + seq_gate = mx.concatenate( + [self.buf_gate[i : i + 1, :r], gate[i : i + 1, :k]], + axis=1) + else: + seq_kv = kv[i : i + 1, :k] + seq_gate = gate[i : i + 1, :k] + if w > 0: + self._pool_lengths[i] += w + if ratio == 4: + if self._prev_kv is None: + # First-ever completion: other rows keep the + # first-window masking (zero kv, -inf gate). + self._prev_kv = mx.zeros( + (B, ratio, kv.shape[2]), dtype=kv.dtype) + self._prev_gate = mx.full( + (B, ratio, gate.shape[2]), -mx.inf, + dtype=gate.dtype) + # The lookback window is the LAST window the replay + # leaves completed. + self._prev_kv[i] = seq_kv[0, (w - 1) * ratio : w * ratio] + self._prev_gate[i] = seq_gate[ + 0, (w - 1) * ratio : w * ratio] + if new_rem > 0: + self.buf_kv[i, :new_rem] = seq_kv[0, -new_rem:] + self.buf_gate[i, :new_rem] = seq_gate[0, -new_rem:] + self.remainder[i] = new_rem + self._processed[i] += k return n def size(self): diff --git a/gmlx/deepseek_v4_mtp.py b/gmlx/deepseek_v4_mtp.py index e4afcf5..0226110 100644 --- a/gmlx/deepseek_v4_mtp.py +++ b/gmlx/deepseek_v4_mtp.py @@ -109,8 +109,6 @@ def __init__(self, config): super().__init__(config) # The verify rollback path needs the rotating undo log on the class. ensure_rollback_attached() - # Hard-disable the L1 shared-APC tier for V4 v1 (spec_engine reads it). - self._kq_apc_mode = None # Set by the DSpark loader: tuple of trunk layer ids whose hc-mean # outputs the drafter consumes. When set, every engine-facing hidden # is PACKED 3D: [raw_4d.flatten(hc) | cap_l0 | cap_l1 | ...]. diff --git a/gmlx/fresh_gate.py b/gmlx/fresh_gate.py new file mode 100644 index 0000000..9349e82 --- /dev/null +++ b/gmlx/fresh_gate.py @@ -0,0 +1,235 @@ +"""Cache-freshness admission gate (GMLX_APC_FRESH_WAIT_MS). + +Sibling fan-out requests that arrive together co-admit into one mixed +prompt batch. Each row's APC pick runs at batch formation, before any +row stores, so every sibling prefills the shared prefix cold and the +batch does that work once per row. The other freshness windows are +already closed: checkpoint-tier and pooling models form prompt batches +one request at a time, and every store commits on the engine thread +before the finish response leaves the generator, so a request fired on +another's completion always sees its records. + +The gate sits on ``BatchGenerator._next``, the admit_gate pattern: on a +tick that would form a prompt batch, it walks the candidates in FCFS +order and cuts the list before the first follower whose shared prefix +with an earlier candidate is not yet covered by any tier (block chain, +exact entries, anchor LRU -- a stats-neutral peek of the same indexes +the pick reads). The held follower keeps its queue position; the next +formation runs after the leader's prefill has published its stores, and +the follower admits warm. Nothing overtakes a held request and running +rows never wait. + +A follower held longer than GMLX_APC_FRESH_WAIT_MS admits cold anyway: +a leader that never covers the shared prefix must not hold its siblings +forever. The timeout is the deadlock kill, not tunable politeness. + +Knobs: + GMLX_APC_FRESH_WAIT_MS hold ceiling in ms (default 500; 0 = off) + GMLX_APC_FRESH_MIN minimum uncovered shared tokens before a + follower is held (default 256) +""" + +from __future__ import annotations + +import logging +import os +import time + +_log = logging.getLogger(__name__) + +_INSTALLED_FLAG = "_kq_gguf_fresh_gate" +_LCP_CHUNK = 4096 + +# Server-wide counters for /v1/metrics. A hold counts once per request +# entering the held state, not per held tick. +_HOLDS = 0 +_LAST_HOLD = "" + + +def fresh_stats() -> dict: + return {"holds": _HOLDS, "last_hold_reason": _LAST_HOLD or None} + + +def _wait_ms() -> float: + try: + return float(os.environ.get("GMLX_APC_FRESH_WAIT_MS", "500")) + except ValueError: + return 500.0 + + +def _hold_min() -> int: + try: + return int(os.environ.get("GMLX_APC_FRESH_MIN", "256")) + except ValueError: + return 256 + + +def _lcp(a, b) -> int: + """Longest common prefix of two token lists. Chunked slice compares + keep the loop in C for deep prompts.""" + n = min(len(a), len(b)) + p = 0 + while p < n: + step = min(_LCP_CHUNK, n - p) + if a[p:p + step] == b[p:p + step]: + p += step + continue + while p < n and a[p] == b[p]: + p += 1 + break + return p + + +def _covered_len(manager, ids, extra_hash: int) -> int: + """Longest stored prefix of ``ids`` across every in-memory tier. + + Reads the same indexes the admission pick reads (block hash chain, + exact entries, anchor LRU) without acquiring blocks, cloning, or + touching hit counters, so a peek never distorts stats or LRU order. + """ + from mlx_vlm import apc as _apc + + tt = tuple(int(t) for t in ids) + best = 0 + with manager.lock: + bs = int(getattr(manager, "block_size", 0) or 0) + table = getattr(manager, "hash_table", None) + if bs > 0 and table: + parent = _apc.SEED_PARENT_HASH + for i in range(len(tt) // bs): + chunk = tt[i * bs:(i + 1) * bs] + h = _apc._hash_tokens(parent, chunk, extra_hash) + blk = table.get(h) + if blk is None or blk.token_ids != chunk: + break + parent = h + best = (i + 1) * bs + for entry in (getattr(manager, "_exact_cache", None) or {}).values(): + p = len(entry.token_ids) + if (entry.extra_hash == extra_hash and best < p < len(tt) + and tt[:p] == tuple(entry.token_ids)): + best = p + for kids, kh in (getattr(manager, "_kq_anchor_cache", None) or {}): + p = len(kids) + if kh == extra_hash and best < p < len(tt) and tt[:p] == kids: + best = p + return best + + +def _note_hold(gen, uid, reason: str) -> bool: + """Stamp the follower's first-held time; True while under the hold + ceiling.""" + held = getattr(gen, "_kq_fresh_held", None) + if held is None: + held = gen._kq_fresh_held = {} + now = time.perf_counter() + first = held.get(uid) + if first is None: + held[uid] = now + global _HOLDS, _LAST_HOLD + _HOLDS += 1 + _LAST_HOLD = reason + _log.info("APC fresh hold: %s", reason) + return True + if (now - first) * 1000.0 > _wait_ms(): + _log.warning( + "APC fresh hold ceiling %.0fms hit: admitting uid=%s cold", + _wait_ms(), uid) + return False + return True + + +def _keep_count(gen): + """Candidates to admit this tick, or None for the full stock cut. + + Runs only when the stock body could form a prompt batch. The head is + never held; the cut lands before the first follower whose uncovered + shared prefix clears the floor, so FCFS order survives truncation. + """ + manager = getattr(gen, "apc_manager", None) + pending = gen._unprocessed_sequences + held = getattr(gen, "_kq_fresh_held", None) + if held: + alive = {s[0] for s in pending} + for uid in [u for u in held if u not in alive]: + del held[uid] + if manager is None or len(pending) < 2 or gen._prompt_batch is not None: + return None + num_to_add = gen.completion_batch_size - len(gen._generation_batch) + if num_to_add < gen.prefill_batch_size: + return None + wait = _wait_ms() + if wait <= 0: + return None + n = min(gen.prefill_batch_size, len(pending)) + if n < 2: + return None + floor = _hold_min() + kept = [list(pending[0][1] or ())] + for i in range(1, n): + ids = list(pending[i][1] or ()) + shared = max((_lcp(ids, k) for k in kept), default=0) + if shared < floor: + kept.append(ids) + continue + extra_hash = gen._apc_extra_hash(pending[i][3] or {}) + covered = _covered_len(manager, ids, extra_hash) + uncovered = shared - covered + if uncovered < floor: + kept.append(ids) + continue + uid = pending[i][0] + reason = (f"uid={uid} shares {shared} tokens with a co-admitted " + f"candidate, {covered} covered; held for the leader's " + f"stores") + if _note_hold(gen, uid, reason): + return i + kept.append(ids) + return None + + +def install_fresh_admission_gate() -> None: + """Hold sibling co-admission until the shared prefix is stored. + + Late-bound monkeypatch on ``BatchGenerator._next``, the admit_gate + pattern: idempotent flag, env kill switch at install, decision + failure degrades to stock admission. Installs after the headroom + gate so a truncated tick still projects memory for the rows it + admits. + """ + from mlx_vlm.generate import ar as _ar + + if getattr(_ar.BatchGenerator._next, _INSTALLED_FLAG, False): + return + if _wait_ms() <= 0: + return + + _orig_next = _ar.BatchGenerator._next + + def _gated_next(self, **kwargs): + try: + k = _keep_count(self) + except Exception: + _log.warning("fresh gate decision failed; admitting", + exc_info=True) + k = None + if k is None: + return _orig_next(self, **kwargs) + pending = self._unprocessed_sequences + kept, tail = pending[:k], pending[k:] + self._unprocessed_sequences = kept + try: + return _orig_next(self, **kwargs) + finally: + # The stock body rebinds the list, and a handler thread may + # append mid-call; splice unconsumed candidates back ahead of + # the held tail, arrivals behind it. Merge, never clobber. + current = self._unprocessed_sequences + kept_ids = {id(s) for s in kept} + leftover = [s for s in current if id(s) in kept_ids] + arrivals = [s for s in current if id(s) not in kept_ids] + self._unprocessed_sequences = leftover + tail + arrivals + + setattr(_gated_next, _INSTALLED_FLAG, True) + _ar.BatchGenerator._next = _gated_next + _log.info("freshness admission gate installed") diff --git a/gmlx/mem_preflight.py b/gmlx/mem_preflight.py new file mode 100644 index 0000000..ec7165b --- /dev/null +++ b/gmlx/mem_preflight.py @@ -0,0 +1,235 @@ +"""Prompt-side memory preflight (GMLX_PREFLIGHT_MEM). + +A request whose KV cannot fit dies mid-stream today, or trips the +pressure machinery late. An error before the SSE stream opens is +retryable and reportable by every client; a mid-stream abort is a +half-answer. This preflight estimates the prompt-side peak (prompt KV +at the model's per-token cost plus the prefill score transient) and +rejects with HTTP 400 and the numbers in the body only on prompt-side +impossibility: the prompt alone cannot fit even with the batch drained. + +Generation-side risk is deliberately not rejected. ``max_tokens`` is +counted only when the client pinned it explicitly, and even then only +prompt plus pinned-max against the drained budget. Default-max requests +are never rejected on generation length; the pressure machinery owns +the tail. Preflight declines the impossible, not the unlikely. + +Every estimate errs on the admit side: quantized KV prices at its bits +without scale overhead, sliding windows cap the token count, MLA prices +the compressed latent, and unprobeable geometry skips the check. Media +requests skip too (image KV and encoder transients are not estimated +in v1). Errors raise a PromptTooLongError subclass, so every existing +handler mapping to 400 applies unchanged. + +Knobs: + GMLX_PREFLIGHT_MEM=0 kill switch, checked per request +""" + +from __future__ import annotations + +import logging +import os + +_log = logging.getLogger(__name__) + +_INSTALLED_FLAG = "_kq_gguf_mem_preflight" +GB = 1e9 + + +_ERR_CLS = None + + +def _preflight_error_cls(): + global _ERR_CLS + if _ERR_CLS is None: + from mlx_vlm.server.generation import PromptTooLongError + + class MemoryPreflightError(PromptTooLongError): + """Prompt-side KV cannot fit the drained working set.""" + + _ERR_CLS = MemoryPreflightError + return _ERR_CLS + + +def _get(cfg, name, default=None): + if cfg is None: + return default + if isinstance(cfg, dict): + return cfg.get(name, default) + return getattr(cfg, name, default) + + +def _lm_config(model): + cfg = getattr(model, "config", None) + text = _get(cfg, "text_config") + return text if text is not None else cfg + + +def kv_layer_costs(model, bytes_per_elem: float = 2.0): + """Per-layer ``(window_or_None, bytes_per_token)`` for the KV cache, + or None when the geometry cannot be read. Admit-side throughout: MLA + prices the compressed latent, a configured sliding window caps every + layer it could apply to.""" + c = _lm_config(model) + layers = _get(c, "num_hidden_layers") + if not isinstance(layers, int) or layers <= 0: + return None + lora = _get(c, "kv_lora_rank") + if isinstance(lora, int) and lora > 0: + rope = _get(c, "qk_rope_head_dim", 0) or 0 + return [(None, (lora + rope) * bytes_per_elem)] * layers + heads = _get(c, "num_attention_heads") + n_kv = _get(c, "num_key_value_heads") or heads + head_dim = _get(c, "head_dim") + hidden = _get(c, "hidden_size") + if not head_dim and heads and hidden: + head_dim = hidden // heads + if not (isinstance(n_kv, int) and n_kv > 0 + and isinstance(head_dim, int) and head_dim > 0): + return None + per_tok = 2 * n_kv * head_dim * bytes_per_elem + window = _get(c, "sliding_window") + window = window if isinstance(window, int) and window > 0 else None + if window is None: + return [(None, per_tok)] * layers + types = _get(c, "layer_types") + if isinstance(types, (list, tuple)) and len(types) == layers: + return [(window if "sliding" in str(t) else None, per_tok) + for t in types] + pattern = _get(c, "sliding_window_pattern") + if isinstance(pattern, int) and pattern > 0: + return [(None if (i + 1) % pattern == 0 else window, per_tok) + for i in range(layers)] + return [(window, per_tok)] * layers + + +def prompt_kv_bytes(costs, tokens: int) -> float: + return sum(bpt * (tokens if w is None else min(tokens, w)) + for w, bpt in costs) + + +def available_drained_bytes(): + """Working set minus zero-copy weights minus the admission reserve: + what a lone request could hold with the batch drained. MLX-tracked + weight allocations are not subtracted, which only admits more.""" + import mlx.core as mx + + from .prefill_decay import untracked_weight_bytes + from .server_memory import admit_reserve_bytes + + try: + ws = float(mx.device_info()["max_recommended_working_set_size"]) + except Exception: + return None + if ws <= 0: + return None + return ws - untracked_weight_bytes() - admit_reserve_bytes(ws) + + +def _need_bytes(model, costs, prompt_tokens: int, gen_tokens: int = 0): + from .prefill_decay import score_transient_bytes + + kv = prompt_kv_bytes(costs, prompt_tokens + gen_tokens) + return kv + score_transient_bytes(model, None, prompt_tokens) + + +def preflight_prompt_memory(rg, prompt, images=None, audio=None, + videos=None, args=None) -> None: + """Raise MemoryPreflightError when the prompt cannot fit. Best + effort: any probe failure admits.""" + try: + if os.environ.get("GMLX_PREFLIGHT_MEM", "1") == "0": + return + if images or audio or videos: + return + model = getattr(rg, "model", None) + if model is None or not isinstance(prompt, str): + return + bits = getattr(rg, "kv_bits", None) + bpe = bits / 8.0 if isinstance(bits, int) and bits > 0 else 2.0 + costs = kv_layer_costs(model, bpe) + if not costs: + return + avail = available_drained_bytes() + if avail is None: + return + # A text token is at least one character, so the character count + # bounds the token count; most requests never pay a tokenize. + if _need_bytes(model, costs, len(prompt)) <= avail: + pinned = _pinned_max_tokens(args) + if not pinned or _need_bytes( + model, costs, len(prompt), pinned) <= avail: + return + from mlx_vlm.server.generation import _count_prompt_tokens + tokens = _count_prompt_tokens(rg._preprocess_request(prompt)) + if tokens <= 0: + return + need = _need_bytes(model, costs, tokens) + pinned = _pinned_max_tokens(args) + need_pinned = (_need_bytes(model, costs, tokens, pinned) + if pinned else 0.0) + if need <= avail and need_pinned <= avail: + return + what = ("prompt" if need > avail + else f"prompt + max_tokens {pinned}") + worst = max(need, need_pinned) + raise _preflight_error_cls()( + f"request cannot fit: {what} needs an estimated " + f"{worst / GB:.1f} GB of KV and prefill transient, but only " + f"{avail / GB:.1f} GB remains with the batch drained " + f"(prompt_tokens={tokens}). Reduce the prompt" + + (" or max_tokens" if need_pinned > avail else "") + ".") + except Exception as e: + from mlx_vlm.server.generation import PromptTooLongError + if isinstance(e, PromptTooLongError): + raise + _log.warning("memory preflight failed; admitting", exc_info=True) + + +def _pinned_max_tokens(args): + """The request's max_tokens, only when the client pinned it away + from the server default.""" + mt = getattr(args, "max_tokens", None) + if not isinstance(mt, int) or mt <= 0: + return 0 + try: + from mlx_vlm.server.generation import get_server_max_tokens + if mt == get_server_max_tokens(): + return 0 + except Exception: + return 0 + return mt + + +def install_memory_preflight() -> None: + """Run the memory preflight on both request entry points. + + ``validate_context_budget`` covers streaming routes before the SSE + stream opens; ``generate`` covers non-streaming handlers and the + gmlx completions route, which call it before any response starts. + Idempotent.""" + from mlx_vlm.server.generation import ResponseGenerator + + if getattr(ResponseGenerator.generate, _INSTALLED_FLAG, False): + return + + _orig_generate = ResponseGenerator.generate + _orig_validate = ResponseGenerator.validate_context_budget + + def _generate(self, prompt, images=None, audio=None, args=None, + videos=None): + preflight_prompt_memory(self, prompt, images, audio, videos, args) + return _orig_generate(self, prompt, images=images, audio=audio, + args=args, videos=videos) + + def _validate(self, prompt, images=None, audio=None, args=None, + videos=None): + _orig_validate(self, prompt, images=images, audio=audio, + args=args, videos=videos) + preflight_prompt_memory(self, prompt, images, audio, videos, args) + + _generate.__dict__[_INSTALLED_FLAG] = True + _validate.__dict__[_INSTALLED_FLAG] = True + ResponseGenerator.generate = _generate + ResponseGenerator.validate_context_budget = _validate + _log.info("memory preflight installed") diff --git a/gmlx/prefix_cache.py b/gmlx/prefix_cache.py index fca5654..a7584b2 100644 --- a/gmlx/prefix_cache.py +++ b/gmlx/prefix_cache.py @@ -13,6 +13,7 @@ from __future__ import annotations import logging +import weakref from collections import OrderedDict from typing import Any @@ -20,9 +21,40 @@ _log = logging.getLogger(__name__) +# Module-wide counters, merged into /v1/cache/stats: spec-path reuse is +# invisible to the APC manager's counters, and a warm turn served from +# this cache must not read as a cold miss (one live model per server). +_HITS = 0 +_HIT_TOKENS = 0 +_STORES = 0 +_LIVE_CACHES: "weakref.WeakSet[SpecPrefixCache]" = weakref.WeakSet() + + +def spec_prefix_stats() -> dict: + return { + "spec_prefix_hits": _HITS, + "spec_prefix_hit_tokens": _HIT_TOKENS, + "spec_prefix_stores": _STORES, + } + + +def spec_prefix_stats_clear() -> None: + global _HITS, _HIT_TOKENS, _STORES + _HITS = 0 + _HIT_TOKENS = 0 + _STORES = 0 + + +def clear_all_spec_prefix_caches() -> None: + """Drop every live cache's entries (the /v1/cache/reset memory + contract: pinned snapshots run to GBs and must not survive it).""" + for cache in list(_LIVE_CACHES): + cache.clear() + _CACHELIST_TAG = "_CacheList" _ARRAYS_TAG = "_ArraysCache" _POOLING_TAG = "_PoolingCache" +_QUANT_TAG = "_QuantizedKV" def _trim_rotating_state(c: Any, keys: mx.array, values: mx.array): @@ -91,7 +123,16 @@ def _snapshot_entry(c: Any) -> Any: if isinstance(c, cache_types("ArraysCache")): return (_ARRAYS_TAG, list(c.cache), c.left_padding, c.lengths) state = c.state - offset = getattr(c, "offset", state[0].shape[2]) + if isinstance(state[0], (tuple, list)): + # Quantized KV (spec KV_BITS path): a (packed, scales, biases) + # triple per side, trimmed to offset by the state property. Group + # quantization runs along head_dim, so the seq axis slices freely. + keys = tuple(mx.contiguous(a) for a in state[0]) + values = tuple(mx.contiguous(a) for a in state[1]) + return (_QUANT_TAG, keys, values, int(c.offset)) + offset = getattr(c, "offset", None) + if offset is None: + offset = state[0].shape[2] _idx = getattr(c, "_idx", 0) if isinstance(c, cache_types("RotatingKVCache")): # When no trim/reorder applies, the returned arrays are the live ring @@ -123,6 +164,12 @@ def _restore_entry(c: Any, snap: Any) -> None: c.left_padding = left_padding c.lengths = lengths return + if isinstance(snap, tuple) and snap[0] == _QUANT_TAG: + _, keys, values, offset = snap + c.keys = tuple(mx.contiguous(a) for a in keys) + c.values = tuple(mx.contiguous(a) for a in values) + c.offset = offset + return keys, values, offset, _idx = snap c.keys, c.values = _owned_pair(keys, values) c.offset = offset @@ -155,6 +202,9 @@ def _collect_snapshot_arrays(snaps: list[Any], out: list[mx.array]) -> None: out.append(left_padding) if lengths is not None: out.append(lengths) + elif isinstance(snap, tuple) and len(snap) == 4 and snap[0] == _QUANT_TAG: + out.extend(snap[1]) + out.extend(snap[2]) elif isinstance(snap, tuple) and len(snap) == 4: out.extend([snap[0], snap[1]]) @@ -220,6 +270,7 @@ def __init__(self, max_entries: int = 4, min_prefix: int = 32, self._min_prefix = min_prefix self._max_bytes = max_bytes self._total_bytes = 0 + _LIVE_CACHES.add(self) def lookup( self, token_ids: mx.array @@ -248,6 +299,9 @@ def lookup( if best is not None and best_len >= self._min_prefix: self._entries.move_to_end(best[1].token_ids) + global _HITS, _HIT_TOKENS + _HITS += 1 + _HIT_TOKENS += best_len return best return None @@ -294,6 +348,8 @@ def store( self._entries[ids] = entry self._entries.move_to_end(ids) self._total_bytes += entry.nbytes + global _STORES + _STORES += 1 while self._entries and ( len(self._entries) > self._max diff --git a/gmlx/queue_cap.py b/gmlx/queue_cap.py new file mode 100644 index 0000000..31cac9e --- /dev/null +++ b/gmlx/queue_cap.py @@ -0,0 +1,190 @@ +"""Queue depth cap with Retry-After (GMLX_QUEUE_DEPTH_CAP). + +The token queue is bounded only by its timeout. Under a fan-out burst +beyond capacity, excess requests hold sockets for up to half an hour +with no tokens and race client-side timeouts. Harness SDKs handle a 503 +with Retry-After by backing off; they handle silent stalls badly. + +The cap rejects before enqueue: a generation request that would push the +waiting queue past GMLX_QUEUE_DEPTH_CAP gets an immediate HTTP 503 with +a JSON body naming the cap and current depth, plus a Retry-After header. +The header value is the estimated drain time (queue depth x recent mean +tokens per request / current aggregate decode rate), clamped to [2, 60] +seconds, static 5 when no stats exist yet. Estimation reads existing +metrics counters only; nothing new lands in the token loop. + +Depth is the engine's waiting census: requests sitting in the server's +request queue plus prompt candidates the batch generator has not yet +admitted. The metrics in-flight gauge is not a usable source: the +gmlx-owned chat routes never call begin_request, so it reads zero under +any load (the depth e2e queue-cap phase caught exactly this). A tiny +publisher wrapper on ``BatchGenerator._next`` keeps a weakref to the +live generator so the check can read its pending list. + +Knobs: + GMLX_QUEUE_DEPTH_CAP waiting-queue cap (default 2 x decode + concurrency; 0 = off) +""" + +from __future__ import annotations + +import importlib +import logging +import os +import weakref + +_log = logging.getLogger(__name__) + +_CAP_FLAG = "_kq_gguf_queue_cap" +_PUB_FLAG = "_kq_gguf_queue_census" +_RETRY_MIN_S = 2 +_RETRY_MAX_S = 60 +_RETRY_DEFAULT_S = 5 + +# Server-wide counters for /v1/metrics. +_REJECTIONS = 0 +_LAST_REJECT = "" + +# Weakref to the live BatchGenerator, published by the census wrapper. +_GEN_REF = None + + +def queue_cap_stats() -> dict: + return {"rejections": _REJECTIONS, + "last_reject_reason": _LAST_REJECT or None} + + +def _decode_concurrency() -> int: + try: + from .decode_batch import decode_batch + return int(decode_batch()) + except Exception: + return 8 + + +def _cap() -> int: + raw = os.environ.get("GMLX_QUEUE_DEPTH_CAP", "") + if raw: + try: + return int(raw) + except ValueError: + pass + return 2 * _decode_concurrency() + + +def _waiting_depth(rg) -> int: + """Requests waiting for a decode slot: the server request queue plus + the batch generator's unadmitted prompt candidates. Both reads are + racy by design; the cap needs magnitude, not a barrier.""" + depth = 0 + qsize = getattr(getattr(rg, "requests", None), "qsize", None) + if callable(qsize): + try: + depth += max(0, int(qsize())) + except Exception: + pass + gen = _GEN_REF() if _GEN_REF is not None else None + pending = getattr(gen, "_unprocessed_sequences", None) + if pending is not None: + depth += len(pending) + return depth + + +def _retry_after_s(metrics, depth: int) -> int: + done = int(getattr(metrics, "_requests_completed", 0) or 0) + toks = int(getattr(metrics, "_completion_tokens_total", 0) or 0) + gen_toks = int(getattr(metrics, "_generated_tokens_total", 0) or 0) + decode_s = float(getattr(metrics, "_decode_time_total_s", 0.0) or 0.0) + if done <= 0 or toks <= 0 or gen_toks <= 0 or decode_s <= 0: + return _RETRY_DEFAULT_S + mean_tokens = toks / done + rate = gen_toks / decode_s + est = depth * mean_tokens / rate + return int(min(max(est, _RETRY_MIN_S), _RETRY_MAX_S)) + + +def check_queue_depth(): + """Return a 503 JSONResponse when the waiting queue is at the cap, + else None. Read-only; a probe failure admits.""" + try: + cap = _cap() + if cap <= 0: + return None + runtime = importlib.import_module("mlx_vlm.server.runtime").runtime + rg = getattr(runtime, "response_generator", None) + if rg is None: + return None + depth = _waiting_depth(rg) + if depth < cap: + return None + retry = _retry_after_s(getattr(runtime, "metrics", None), depth) + global _REJECTIONS, _LAST_REJECT + _REJECTIONS += 1 + _LAST_REJECT = f"depth {depth} at cap {cap}, retry {retry}s" + _log.info("queue cap: rejected request (%s)", _LAST_REJECT) + from fastapi.responses import JSONResponse + return JSONResponse( + status_code=503, + content={"error": { + "message": (f"server queue is full: {depth} requests " + f"waiting, cap {cap}; retry after {retry}s"), + "type": "server_overloaded", + "queue_depth": depth, + "queue_cap": cap, + }}, + headers={"Retry-After": str(retry)}, + ) + except Exception: + _log.warning("queue cap probe failed; admitting", exc_info=True) + return None + + +def _install_census() -> None: + """Publish a weakref to the live BatchGenerator from its tick. + + Pure passthrough wrapper; the engine swaps generators across idle + gaps, so the ref re-publishes whenever the instance changes.""" + from mlx_vlm.generate import ar as _ar + + if getattr(_ar.BatchGenerator._next, _PUB_FLAG, False): + return + _orig_next = _ar.BatchGenerator._next + + def _published_next(self, **kwargs): + global _GEN_REF + ref = _GEN_REF + if ref is None or ref() is not self: + _GEN_REF = weakref.ref(self) + return _orig_next(self, **kwargs) + + setattr(_published_next, _PUB_FLAG, True) + _ar.BatchGenerator._next = _published_next + + +def install_queue_depth_cap() -> None: + """Wrap the generation POST routes with the queue-depth check. + + The check runs before the stock handler, so a rejected request never + opens an SSE stream and never reaches the engine. Idempotent per + route; GMLX_QUEUE_DEPTH_CAP=0 disables at install. + """ + if _cap() <= 0: + return + _install_census() + from .server_patches._common import _CHAT_PATHS, _wrap_post_routes + + app = importlib.import_module("mlx_vlm.server.app").app + paths = _CHAT_PATHS + ("/responses", "/v1/responses", + "/messages", "/v1/messages", + "/completions", "/v1/completions") + + def _make(original): + async def endpoint(*args, **kwargs): + rejected = check_queue_depth() + if rejected is not None: + return rejected + return await original(*args, **kwargs) + return endpoint + + _wrap_post_routes(app, paths, _CAP_FLAG, _make) + _log.info("queue depth cap installed (cap=%d)", _cap()) diff --git a/gmlx/seed_rows.py b/gmlx/seed_rows.py new file mode 100644 index 0000000..a091678 --- /dev/null +++ b/gmlx/seed_rows.py @@ -0,0 +1,123 @@ +"""Per-request seed under batching. + +Eval harnesses and agent frameworks send ``seed`` per request. The batch +engine builds one sampler when the generator spins up, so before this, +only the first request's seed took effect and it colored every row. +Now each request's seed rides with its row: the shared sampler keeps a +uid-to-seed registry, and every keyed draw derives that row's key from +its own seed. Rows without a seed keep the stock derivation byte for +byte, so seeded and unseeded rows coexist in one batch. + +Honest semantics (also in docs/server-config.md): seed guarantees a +deterministic sampling stream for that request. It does not guarantee +bitwise-identical output across runs with different batch composition, +because batched matmul reduction order shifts logits at float +tolerance. Same composition (for example solo replay) reproduces +exactly. With speculation, drafts are greedy and the target draws for a +seeded B=1 request come from the same per-request key stream, so a +same-setting replay matches; replays across different speculation +settings do not. + +Wiring: the request's seed is stashed when the engine builds the +request's per-row hooks (the last per-request step before insert on the +single GPU thread) and bound to the uid insert returns. The decode step +and the speculative round both publish their row uids on the sampler +around each draw. +""" + +from __future__ import annotations + +import logging + +_log = logging.getLogger(__name__) + +_INSTALLED_FLAG = "_kq_gguf_seed_rows" +_MAX_SEEDS = 1024 + +# Single-slot handoff between the per-request argument hook and the +# insert that follows it on the GPU thread. +_PENDING: list = [] + + +def register_row_seed(sampler, uid, seed) -> None: + seeds = getattr(sampler, "_kq_row_seeds", None) + if seeds is None: + return + seeds[uid] = int(seed) + while len(seeds) > _MAX_SEEDS: + seeds.pop(next(iter(seeds))) + + +def install_per_request_seed() -> None: + """Bind each request's seed to its batch row. Idempotent.""" + from mlx_vlm.generate import ar as _ar + from mlx_vlm.server.generation import ResponseGenerator + + if getattr(_ar.BatchGenerator.insert, _INSTALLED_FLAG, False): + return + + _orig_criteria = ResponseGenerator._make_thinking_budget_criteria + + def _criteria_with_seed(self, args, input_ids): + _PENDING.clear() + seed = getattr(args, "seed", None) + if seed is not None and getattr(args, "temperature", 1.0) != 0: + _PENDING.append(int(seed)) + return _orig_criteria(self, args, input_ids) + + _orig_insert = _ar.BatchGenerator.insert + + def _insert_with_seed(self, *args, **kwargs): + seed = _PENDING.pop() if _PENDING else None + uids = _orig_insert(self, *args, **kwargs) + if seed is not None and getattr(self, "sampler", None) is not None: + for uid in uids: + register_row_seed(self.sampler, uid, seed) + return uids + + _orig_step = _ar.GenerationBatch._step + + def _step_with_rows(self): + sampler = self.sampler + if getattr(sampler, "_kq_row_seeds", None): + sampler._kq_rows = list(self.uids) + try: + return _orig_step(self) + finally: + sampler._kq_rows = None + return _orig_step(self) + + _orig_generate = _ar.PromptProcessingBatch.generate + + def _generate_with_rows(self, sampler, *args, **kwargs): + if getattr(sampler, "_kq_row_seeds", None): + sampler._kq_rows = list(self.uids) + try: + return _orig_generate(self, sampler, *args, **kwargs) + finally: + sampler._kq_rows = None + return _orig_generate(self, sampler, *args, **kwargs) + + _orig_next = _ar.SpeculativeGenerationBatch.next + + def _next_with_rows(self): + sampler = getattr(self, "sampler", None) + if getattr(sampler, "_kq_row_seeds", None): + sampler._kq_rows = list(self._all_uids) + try: + return _orig_next(self) + finally: + sampler._kq_rows = None + return _orig_next(self) + + _criteria_with_seed.__dict__[_INSTALLED_FLAG] = True + _insert_with_seed.__dict__[_INSTALLED_FLAG] = True + _step_with_rows.__dict__[_INSTALLED_FLAG] = True + _generate_with_rows.__dict__[_INSTALLED_FLAG] = True + _next_with_rows.__dict__[_INSTALLED_FLAG] = True + ResponseGenerator._make_thinking_budget_criteria = _criteria_with_seed + _ar.BatchGenerator.insert = _insert_with_seed + _ar.GenerationBatch._step = _step_with_rows + _ar.PromptProcessingBatch.generate = _generate_with_rows + _ar.SpeculativeGenerationBatch.next = _next_with_rows + _log.info("per-request seed installed") diff --git a/gmlx/server_patches/__init__.py b/gmlx/server_patches/__init__.py index 12fd236..b43b2b9 100644 --- a/gmlx/server_patches/__init__.py +++ b/gmlx/server_patches/__init__.py @@ -186,6 +186,8 @@ def install_server_patches(cfg, *, reload_fn=None) -> None: install_step_timing() if os.environ.get("GMLX_DISABLE_FAST_SAMPLER") != "1": install_fast_sampler() + from ..seed_rows import install_per_request_seed + install_per_request_seed() from .. import spec_engine spec_engine.install_full_prompt_mtp_prefill() spec_engine.install_owned_spec_engine() @@ -212,6 +214,10 @@ def install_server_patches(cfg, *, reload_fn=None) -> None: # no prefill work, and decode runs unpaced while admission waits. from ..admit_gate import install_admit_headroom_gate install_admit_headroom_gate() + # After the headroom gate so a truncated tick still projects memory + # for the rows it admits. + from ..fresh_gate import install_fresh_admission_gate + install_fresh_admission_gate() install_chat_template_kwargs() install_thinking_budget_fix() install_stream_timings() @@ -254,6 +260,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() + from ..queue_cap import install_queue_depth_cap + install_queue_depth_cap() + from ..mem_preflight import install_memory_preflight + install_memory_preflight() # 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 diff --git a/gmlx/server_patches/routes.py b/gmlx/server_patches/routes.py index f8e1b11..c5d06e8 100644 --- a/gmlx/server_patches/routes.py +++ b/gmlx/server_patches/routes.py @@ -256,6 +256,18 @@ def snapshot(): base["admission"] = admit_stats() except Exception: pass + try: + from ..fresh_gate import fresh_stats + + base["freshness"] = fresh_stats() + except Exception: + pass + try: + from ..queue_cap import queue_cap_stats + + base["queue"] = queue_cap_stats() + except Exception: + pass return base snapshot.__dict__[_PATCH_FLAG] = True diff --git a/gmlx/server_patches/sampling.py b/gmlx/server_patches/sampling.py index b7494eb..b683e7c 100644 --- a/gmlx/server_patches/sampling.py +++ b/gmlx/server_patches/sampling.py @@ -218,6 +218,24 @@ def __init__(self, *, temperature, top_p=1.0, top_k=0, min_p=0.0, seed=None): self.top_k = int(top_k or 0) self.min_p = float(min_p or 0.0) self.seed = DEFAULT_SEED if seed is None else int(seed) + # Per-request seeds (uid -> seed) and the row->uid context for the + # draw in flight, maintained by gmlx.seed_rows. A row with a + # registered seed draws from its own key stream; every other row + # keeps the stock derivation byte for byte. + self._kq_row_seeds: dict = {} + self._kq_rows = None + + def _row_keys(self, row_ids, positions): + import mlx.core as mx + from mlx_vlm.server.generation import _position_keys, _position_seed + rows = self._kq_rows + seeds = self._kq_row_seeds + if not seeds or rows is None or len(rows) != len(row_ids): + return _position_keys(self.seed, row_ids, positions) + return mx.stack([ + mx.random.key(_position_seed( + seeds.get(u, self.seed), row, pos)) + for u, row, pos in zip(rows, row_ids, positions)]) @property def _has_filter(self): @@ -274,10 +292,9 @@ def __call__(self, logprobs): def sample_target(self, logprobs, *, row_ids, positions): import mlx.core as mx - from mlx_vlm.server.generation import _position_keys if logprobs.shape[0] != len(row_ids) or len(row_ids) != len(positions): raise ValueError("row_ids and positions must match logprobs batch size.") - keys = _position_keys(self.seed, row_ids, positions) # [B, 2] + keys = self._row_keys(row_ids, positions) # [B, 2] def _cat(row, key): return mx.random.categorical(row, key=key) diff --git a/gmlx/spec_engine.py b/gmlx/spec_engine.py index 732f04b..c58363f 100644 --- a/gmlx/spec_engine.py +++ b/gmlx/spec_engine.py @@ -137,6 +137,10 @@ def _install_apc_manager_stash() -> None: _orig_init = BatchGenerator.__init__ def _init_with_stash(self, model, processor, **kwargs): + # upstream server never passes completion_batch_size; inject ours + if "completion_batch_size" not in kwargs: + from .decode_batch import decode_batch + kwargs["completion_batch_size"] = decode_batch() # Kill switch (re-read per call): with spec APC off, stock ar.py # must not see the manager on the speculative path either -- since # mlx-vlm 0.6.4 its own post-prefill exact store handles B=1 MTP @@ -149,6 +153,21 @@ def _init_with_stash(self, model, processor, **kwargs): except Exception: pass _orig_init(self, model, processor, **kwargs) + # APC arrived armed but upstream's quantized-KV opt-out dropped it + # (ar.py nulls the manager whenever kv_bits is set; no tier serves + # quantized caches). The mode probe still reads "block" for these + # models, so without this line the server boots silent and every + # request prefills cold. Draft-model batches are excluded: upstream + # nulls their manager by design and the owned ladder resolves (and + # warns) through _resolve_l1. + if (kwargs.get("apc_manager") is not None + and kwargs.get("kv_bits") is not None + and kwargs.get("draft_model") is None + and getattr(self, "apc_manager", None) is None): + _log.warning( + "APC OFF: KV quantization (kv_bits=%s) opts out of the " + "block APC tier upstream -- every request prefills cold", + kwargs.get("kv_bits")) # Ckpt-tier models form prompt batches one request at a time: the # owned APC declines B>1 prefill, so a coalesced burst would go # all-cold, and B>1 prompt batching is not a throughput win anyway @@ -290,6 +309,13 @@ def _l1_lookup_and_arm_store(batch, manager, mode, l0_prefix) -> int: prefix_len = 0 tier = "exact" pick = view._apc_pick_for((0, ids_list, 0, prompt_kwargs, None, None)) + # Same trivial-pick floor as the admission wrapper: a sub-block + # exact restore saves nothing and its nonzero l1_prefix would skip + # the L0 hidden store for this request. + if (pick is not None and not pick.get("matched_blocks") + and 0 < int(pick.get("prefix_len") or 0) + < int(manager.block_size)): + pick = None if pick is not None: warm = pick.get("warm_cache") blocks = list(pick.get("matched_blocks") or ()) @@ -334,6 +360,15 @@ def _l1_lookup_and_arm_store(batch, manager, mode, l0_prefix) -> int: manager.release(blocks) blocks = [] warm, prefix_len, tier = aw, ap, "anchor" + if warm and 0 < prefix_len < len(ids_list) and tier in ( + "exact", "anchor"): + # Same batch-aware merge admission applies to its picks: raw + # exact/anchor clones carry single-row leaves (left_padding + # None, scalar offsets) and crash the batch cache classes' + # update path (mx.depends on a None) when the suffix forwards. + from mlx_vlm import apc as _apc + warm, _ = _apc.make_warm_batch_exact_cache_multi( + [warm], [prefix_len]) if warm and 0 < prefix_len < len(ids_list): batch.prompt_cache = warm # Matched blocks stay acquired until the stock post-prefill @@ -928,6 +963,15 @@ def _pick_with_anchor(self, sequence): _uid, ids_list, _mt, prompt_kwargs, _lps, _crit = sequence if not ids_list or len(ids_list) < 2: return pick + # Floor trivial exact picks: a sub-block restore (a bare-BOS + # match off an unrelated request) saves nothing but suffix- + # constructs the batch, knocking the spec path's ids out of + # render space (anchor + retirement keys). Real warm picks are + # thousands of tokens and pass untouched. + if (pick is not None and not pick.get("matched_blocks") + and 0 < int(pick.get("prefix_len") or 0) + < int(manager.block_size)): + pick = None have = int((pick or {}).get("prefix_len") or 0) extra_hash = self._apc_extra_hash(prompt_kwargs or {}) floor = max(have, self._apc_safe_prefix_lookup_min(ids_list)) @@ -1079,6 +1123,7 @@ def _mtp_prefill_init(batch) -> None: batch._mtp_l1_prefix_len = 0 if batch._inputs_embeds is None: + _log.info("KQDBG mtp_prefill_init: inputs_embeds None, ladder skipped") return # Gated to B=1 because PromptProcessingBatch prefills one request at a @@ -1097,6 +1142,18 @@ def _mtp_prefill_init(batch) -> None: "(owned-path APC requires single-request prefill)", b) return + # Upstream admission already restored a prefix and built this batch + # suffix-only: the owned ladder's keys (L0 and L1 both) are full-prompt + # token ids, so every lookup and store here would run in the wrong + # space -- a suffix-keyed L0 entry cross-hits a later turn's suffix and + # its restore clobbers the upstream warm cache. Leave these batches to + # the stock machinery, which owns their meta and store schedule. + up_meta = getattr(batch, "_apc_meta", None) or [] + if up_meta and isinstance(up_meta[0], dict) \ + and int(up_meta[0].get("prefix_len") or 0) > 0: + batch._mtp_upstream_warm = True + return + restored = 0 spec_cache = _get_spec_prefix_cache(batch.model) if spec_cache is not None: @@ -1423,7 +1480,8 @@ def _mtp_generate(self, sampler, stop_criteria, b = int(full_hidden.shape[0]) if full_ids is not None else 0 spec_cache = ( _get_spec_prefix_cache(self.model) - if b == 1 and l1_prefix == 0 else None + if b == 1 and l1_prefix == 0 + and not getattr(self, "_mtp_upstream_warm", False) else None ) if spec_cache is not None and full_ids is not None: spec_cache.store(full_ids, result.prompt_cache, full_hidden) @@ -1610,12 +1668,15 @@ def _lift_host_cache(c): def _preempt_scalar(self) -> bool: """Preempt a live scalar (B=1) spec generation so queued rows can - join: close the generator at its round boundary (its GeneratorExit - handler rolls the target cache back to the delivered tokens), lift - the caches to batch classes, and mark the batch armless - (hidden=None); _start_rounds then rebuilds it on the batch loop, - whose first injection drain admits the waiters. GMLX_MTP_PREEMPT=0 - leaves the old drain-wait behavior. + join: close the generator, deliver the closed round's undelivered + tail (the scalar path yields one token per next(), so a close + usually lands mid-round; those tokens are verified and their KV + stays in the cache), lift the caches to batch classes, and mark + the batch armless (hidden=None); _start_rounds then rebuilds it on + the batch loop, whose first injection drain admits the waiters. + The rebuild resumes from the round's bonus token, whose KV is not + in the cache. GMLX_MTP_PREEMPT=0 leaves the old drain-wait + behavior. The rebuilt row carries no APC retirement context (batch-loop rows start with retire_ctxs None), so the preempted request's prefix is @@ -1628,9 +1689,36 @@ def _preempt_scalar(self) -> bool: if last is None: return False it = self._rounds_iter + captured = [] if it is not None: self._rounds_iter = None - it.close() + self.model._kq_preempt_capture = captured + try: + it.close() + finally: + try: + del self.model._kq_preempt_capture + except AttributeError: + pass + responses = [] + uid = self._all_uids[0] + for tok in captured: + if self._finished[0]: + break + tok = int(tok) + self._num_tokens[0] += 1 + finish = self._finish_reason(0, tok) + if finish is not None: + self._finished[0] = True + responses.append(self.Response( + uid=uid, token=tok, token_logprob=0.0, finish_reason=finish)) + last = tok + self._kq_preempt_responses = responses + if self._finished[0]: + # The captured tail finished the row; nothing to rebuild. The + # pending injections promote through __len__ once drained. + self._refresh_uids() + return False self.prompt_cache = [_lift_host_cache(c) for c in self.prompt_cache] self.first_tokens = mx.array([int(last)], dtype=self.token_dtype) self.hidden = None @@ -1653,8 +1741,11 @@ def _next_with_injection(self): preempted = False if pending and len(self._all_uids) == 1: preempted = _preempt_scalar(self) + # The preempt capture: verified tokens the closed round had not yet + # delivered. They precede everything this call returns. + pre_responses = self.__dict__.pop("_kq_preempt_responses", None) or [] if pending and (len(self._all_uids) > 1 or preempted): - responses = [] + responses = list(pre_responses) gen_inj = getattr(self.model, "_generator_injections", None) if gen_inj is None: self.model._generator_injections = [] @@ -1703,7 +1794,7 @@ def _next_with_injection(self): _release_if_finished(self) return responses - responses = _orig_next(self) + responses = pre_responses + _orig_next(self) _note_last_tokens(self, responses) _release_if_finished(self) return responses diff --git a/gmlx/speculative.py b/gmlx/speculative.py index 171928b..37be22b 100644 --- a/gmlx/speculative.py +++ b/gmlx/speculative.py @@ -524,8 +524,29 @@ def _stochastic_walk(lm, verify, draft_tokens: mx.array, sampler, budget: int, ) +def _seeded_target_draw(sampler, logprobs, base_pos): + """Target draws for a seeded lone row: per-position keys from the + request's own seed, through the sampler's keyed path. Falls back to + the process stream (stock behavior) whenever the row is not seeded + or the sampler has no keyed path.""" + rows = getattr(sampler, "_kq_rows", None) + seeds = getattr(sampler, "_kq_row_seeds", None) + target = getattr(sampler, "sample_target", None) + if (base_pos is None or not seeds or not rows or target is None + or len(set(rows)) != 1 or seeds.get(rows[0]) is None): + return sampler(logprobs) + n = int(logprobs.shape[0]) + saved = sampler._kq_rows + sampler._kq_rows = [rows[0]] * n + try: + return target(logprobs, row_ids=[0] * n, + positions=[base_pos + 1 + j for j in range(n)]) + finally: + sampler._kq_rows = saved + + def _coupled_walk(lm, verify, draft_tokens: mx.array, sampler, budget: int, - top2=None, pq=None): + top2=None, pq=None, base_pos=None): """Rejection walk with a single host sync. Sample every verify position into one deferred graph (sequentially, so the @@ -563,7 +584,8 @@ def _coupled_walk(lm, verify, draft_tokens: mx.array, sampler, budget: int, if sampler is None: target = mx.argmax(logprobs, axis=-1) # [n_pos] else: - target = sampler(logprobs).reshape(-1) # [n_pos] + target = _seeded_target_draw( + sampler, logprobs, base_pos).reshape(-1) # [n_pos] draft_row = draft_tokens.reshape(-1) if n_draft > 0: match = (target[:n_draft] == draft_row).astype(mx.int32) @@ -840,7 +862,8 @@ def draft_sampler(logits, _inner=_pq_inner_sampler, _stash=pq_stash): lm, verify, draft_tokens, _walk_sampler, max_tokens - emitted, top2=list(top2_stash) if top2_stash else None, - pq=list(pq_stash) if pq_stash else None) + pq=list(pq_stash) if pq_stash else None, + base_pos=emitted) stoch_stash.clear() if top2_stash is not None: # Consumed; the accept hook below re-seeds entry 0 for the @@ -863,6 +886,24 @@ def draft_sampler(logits, _inner=_pq_inner_sampler, _stash=pq_stash): delivered += 1 yield tok except GeneratorExit: + capture = getattr(model, "_kq_preempt_capture", None) + if capture is not None: + # Preempt close, not a stop: the round completed and + # verified, so the undelivered tail is real output. + # Hand it to the rebuilder instead of trimming it -- + # a trim here loses one token per preempt, because the + # rebuild re-forwards the last DELIVERED token whose + # KV the trim left in the cache (duplicate position) + # and the first undelivered token never reaches the + # wire. Keep the accepts (normal round-end state); + # the bonus stays out of the cache as the rebuild's + # first input. + capture.extend(int(t) for t in new_tokens[delivered:]) + if _has_rollback and accepted < bs - 1: + with mx.stream(generation_stream): + _rollback_fn(prompt_cache, verify.gdn_states, + accepted, bs) + raise # Consumer stopped mid-round (EOS / stop string). Roll the target # cache back to exactly the delivered tokens so the finish seam # sees KV consistent with what was consumed (APC retirement @@ -1122,6 +1163,12 @@ def owned_server_rounds( # the locals, so this can't double-store. See _retire_b1. if rounds is not None: rounds.close() + if (retire_ctx is not None + and getattr(model, "_kq_preempt_capture", None) is not None): + # Preempt close: the cache keeps the captured round tail past + # `generated`, so a snapshot here would disagree with the + # delivered stream. The rebuilt row drops retirement by design. + retire_ctx = None if retire_ctx is not None: _retire_b1(model, prompt_cache, generated, retire_ctx, drafter=drafter, sidecar_ctx=sidecar_ctx) @@ -1678,6 +1725,15 @@ def _reset_armed(n: int) -> None: drafter.reset(model, left_padding=[0] * n) except (TypeError, ValueError): drafter.reset(model) + except NotImplementedError: + # B=1-only drafters refuse any padding list, even the trivial + # [0] a width-1 batch passes (a preempted scalar rebuilds into + # this loop at B=1; formation gates wider batches and + # injections trip the gate before crossing the cap). One row + # has no left padding, so the bare scalar arm is identical. + if n != 1: + raise + drafter.reset(model) # hidden=None with the gate open means no prefill capture exists (a # preempted scalar generation rebuilt into this loop): the first round @@ -1744,8 +1800,10 @@ def _reset_armed(n: int) -> None: # Double buffer for gated (plain-decode) rounds: the tokens the GPU is # already computing. None means "re-prime", set whenever the batch shape - # changes under us (injection, row retirement, adoption). + # changes under us (row retirement, adoption). An injection instead + # holds admission one round so the buffer is consumed, never discarded. _gated_pending = None + _inject_hold = False def _gated_step(inputs): """One plain target decode step: [n_active] tokens in, [n_active] out. @@ -1813,13 +1871,20 @@ def _arm_capture(): def _drain_injections(): # continuous-batch injection - nonlocal hidden, B_orig, _gated_pending + nonlocal hidden, B_orig, _gated_pending, _inject_hold gen_inj = getattr(model, "_generator_injections", None) if gen_inj: + if _gated_pending is not None: + # The buffer holds a dispatched step whose input KV is + # already in the target cache; it must be consumed, never + # discarded. Discarding skips one delivered token per row, + # and the re-prime re-forwards its inputs, duplicating + # their KV. Hold admission one round: the caller consumes + # the step without re-dispatching, and the next boundary + # admits with the buffer empty. + _inject_hold = True + return n_before = B_orig - # The batch is about to widen; anything already dispatched was - # shaped for the old width, so re-prime rather than slice. - _gated_pending = None # Trip BEFORE processing: a tripping drain must not pay # inject_rows, which teacher-forces the whole injected prompt # through a drafter this batch will never use again. The queue can @@ -1928,7 +1993,8 @@ def _drain_injections(): # dispatched next step whose input KV is already in the cache, so it # must be consumed, never discarded: one more plain round runs # without re-dispatching, and the round after that arms. - dispatch_next = True + dispatch_next = not _inject_hold + _inject_hold = False if gated and _resume_ready(): if _gated_pending is None: gated = False @@ -2080,8 +2146,12 @@ def _drain_injections(): if any(a < bs - 1 for a in accepted_list) and _has_rollback: with mx.stream(generation_stream): - _rollback_fn(prompt_cache, verify.gdn_states, - accepted_list, bs) + # Rollback hooks are scalar-only: every model that defines + # one is B=1-limited, so spec rounds with a rollback only + # run at width 1 here (wider batches gate at formation). + # Pass that row's int, not a one-element list. + acc = accepted_list[0] if n_active == 1 else accepted_list + _rollback_fn(prompt_cache, verify.gdn_states, acc, bs) if _needs_shared_kv and not gated: rejected_global = bs - (max_a + 1) @@ -2152,9 +2222,11 @@ def _drain_injections(): if callable(filter_drafter): filter_drafter(keep_mx) hidden = hidden[keep_mx] - for k in next_shared_kv: - K_next, V_next = next_shared_kv[k] - next_shared_kv[k] = (K_next[keep_mx], V_next[keep_mx]) + if _needs_shared_kv: + for k in next_shared_kv: + K_next, V_next = next_shared_kv[k] + next_shared_kv[k] = ( + K_next[keep_mx], V_next[keep_mx]) active_idx = [active_idx[j] for j in keep_slots] if not gated: diff --git a/tests/e2e/run_apc_depth_e2e.py b/tests/e2e/run_apc_depth_e2e.py index 4d42f3a..4aa479a 100644 --- a/tests/e2e/run_apc_depth_e2e.py +++ b/tests/e2e/run_apc_depth_e2e.py @@ -47,6 +47,18 @@ churn distinct mid-size prefixes cycle records through the ckpt LRU, then the original replay must still serve correctly with no exception-declines + burst (block/exact) K siblings sharing a cold user-turn prefix + fire at once; on the block tier the fresh gate must hold + the followers for the leader's stores so they admit warm, + on the exact tier the numbers are noted (user-turn stores + land at retirement, past the hold ceiling) + queue-cap (only with --queue-cap) the main servers boot with + GMLX_QUEUE_DEPTH_CAP=1; a flood past the decode batch + must draw immediate 503s carrying the Retry-After + contract while the rest of the flood serves, the + rejections counter must move, and a retry after the + drain must succeed (Unit 5's wire contract, which unit + tests fake) session (only with --session N) an agent-shaped conversation on a dedicated server: N turns of growing history with streamed replies, tool-call/tool-role messages, a mid-stream client @@ -291,6 +303,43 @@ def http_chat( return st, text, content, ptok, wall +def raw_chat( + base: str, + mid: str, + messages: list, + *, + max_tokens: int = 24, + timeout: float = 900.0, +): + """One chat POST that surfaces the raw HTTP contract: (status, + headers, parsed JSON body). A rejection returns its error body and + headers instead of raising (queue-cap phase reads Retry-After).""" + body = { + "model": mid, + "messages": messages, + "max_tokens": max_tokens, + "temperature": 0.0, + } + req = urllib.request.Request( + f"{base}/v1/chat/completions", + data=json.dumps(body).encode(), + method="POST", + ) + req.add_header("Content-Type", "application/json") + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: + payload = resp.read().decode("utf-8", "replace") + return resp.status, resp.headers, json.loads(payload or "{}") + except urllib.error.HTTPError as e: + try: + parsed = json.loads(e.read().decode("utf-8", "replace") or "{}") + except json.JSONDecodeError: + parsed = {} + return e.code, e.headers, parsed + except Exception as e: # noqa: BLE001 - report, don't raise + return -1, {}, {"error": {"message": f"{type(e).__name__}: {e}"}} + + def sse_chat( base: str, mid: str, @@ -483,11 +532,35 @@ def main() -> int: "blocks (~1.3 tok/word) shared by every request, instead of the " "one-line default; the shape the system-prompt anchor serves", ) + ap.add_argument( + "--queue-cap", + action="store_true", + help="queue-cap phase: boot the main servers with " + "GMLX_QUEUE_DEPTH_CAP=1 and flood past the decode batch; tail " + "arrivals must draw an immediate 503 with the Retry-After " + "contract and the server must recover", + ) ap.add_argument( "--no-tripwire", action="store_true", help="skip the missed-adoption tripwire phase (boots a third server)", ) + ap.add_argument( + "--reply-tokens", + type=int, + default=SESSION_REPLY_TOKENS, + help="session-phase reply budget (default 512); reasoning-channel " + "models at depth need 2048+ or replies truncate mid-think with " + "empty content and the turn chain degenerates", + ) + ap.add_argument( + "--system-suffix", + default=None, + metavar="TEXT", + help="text appended to every system message, e.g. 'Reasoning " + "strength: low.' for ATEM reasoning models whose default strength " + "meanders past any reply budget at depth", + ) ap.add_argument( "--require-idle", action="store_true", @@ -594,11 +667,12 @@ def content_clean(content: str) -> bool: f"{f2['reading']} units by {f2['technician']}" ) - system_msg = agent_system(a.system_words) if a.system_words > 0 \ - else SYSTEM_MSG + sfx = ("\n\n" + a.system_suffix) if a.system_suffix else "" + system_msg = (agent_system(a.system_words) if a.system_words > 0 + else SYSTEM_MSG) + sfx def mk_msgs(user_content: str, system: str | None = None) -> list: - system = system_msg if system is None else system + system = system_msg if system is None else system + sfx msgs = [] if a.no_system else [{"role": "system", "content": system}] msgs.append({"role": "user", "content": user_content}) return msgs @@ -606,6 +680,14 @@ def mk_msgs(user_content: str, system: str | None = None) -> list: def tick(key: str) -> int: return int(stats(base).get(key, 0) or 0) + def reuse_tick() -> int: + # Total reused prefix tokens: APC adoption plus spec prefix-cache + # restores. The MTP serve path's first reuse layer never moves the + # manager's matched counter, but a warm turn is warm either way. + s = stats(base) + return int(s.get(K["matched"], 0) or 0) + int( + s.get("spec_prefix_hit_tokens", 0) or 0) + def settle(key: str, timeout: float = 15.0) -> int: # retirement stores are async; wait for the counter to hold still last = tick(key) @@ -646,6 +728,10 @@ def settle(key: str, timeout: float = 15.0) -> int: scheme=a.scheme, bits=a.bits_a, ) + if a.queue_cap: + # cap 1 = one waiter past the decode batch. Unreachable by every + # other phase: their concurrency stays far below the batch size. + env_a["GMLX_QUEUE_DEPTH_CAP"] = "1" t_load0 = time.monotonic() with ServerProc( serve_args(a.model, a.speculative, a.draft_gguf), @@ -767,39 +853,49 @@ def settle(key: str, timeout: float = 15.0) -> int: # -- anchor (only with --system-words) ------------------------------ # A shared agent-shaped system prompt makes the divergent request a - # sibling: it must restore the system block from the anchor stored - # during the cold request, on every tier. Without --system-words the - # divergence sits a few tokens in and no anchor arms at all, which - # is why this gate is opt-in rather than always on. + # sibling: on the ckpt and exact tiers it must restore the system + # block from the anchor stored during the cold request; the block + # tier serves the same prefix from the block chain with no anchor. + # Without --system-words the divergence sits a few tokens in and no + # anchor arms at all, which is why this gate is opt-in. if a.system_words > 0 and not a.no_system: st_all = stats(base) - anchored = int(st_all.get("anchor_stores", 0) or 0) + int( - st_all.get("ckpt_stores", 0) or 0 - ) # ~1.3 tok/word, less the template scaffolding and the grid or # chunk snap the ckpt tier applies below the boundary. anchor_floor = int(a.system_words * 0.6) - rep.check( - "anchor.armed", - anchored > 0, - f"anchor_stores={st_all.get('anchor_stores', 0)} " - f"ckpt_stores={st_all.get('ckpt_stores', 0)}", - ) rep.check( "anchor.divergent_adopted", dm >= anchor_floor, f"divergent matched +{dm} >= system-anchor floor " f"{anchor_floor} ({a.system_words} system words)", ) - served = int(st_all.get("anchor_hits", 0) or 0) + int( - st_all.get("ckpt_hits", 0) or 0 - ) - rep.check( - "anchor.served", - served > 0, - f"anchor_hits={st_all.get('anchor_hits', 0)} " - f"ckpt_hits={st_all.get('ckpt_hits', 0)}", - ) + if a.tier == "block": + # Pure-KV models never enter the anchor code: the stock + # block tier already serves siblings at block granularity. + rep.note( + "anchor.scope", + "block tier serves the system prefix from the block " + "chain; no anchor arms", + ) + else: + anchored = int(st_all.get("anchor_stores", 0) or 0) + int( + st_all.get("ckpt_stores", 0) or 0 + ) + rep.check( + "anchor.armed", + anchored > 0, + f"anchor_stores={st_all.get('anchor_stores', 0)} " + f"ckpt_stores={st_all.get('ckpt_stores', 0)}", + ) + served = int(st_all.get("anchor_hits", 0) or 0) + int( + st_all.get("ckpt_hits", 0) or 0 + ) + rep.check( + "anchor.served", + served > 0, + f"anchor_hits={st_all.get('anchor_hits', 0)} " + f"ckpt_hits={st_all.get('ckpt_hits', 0)}", + ) # -- turns --------------------------------------------------------- # A real conversation through the real chat template: the @@ -829,11 +925,11 @@ def history_safe(reply: str) -> str: for uq, at in history: msgs.append({"role": "assistant", "content": at}) msgs.append({"role": "user", "content": uq}) - m0 = tick(K["matched"]) + m0 = reuse_tick() st, turn_text, last_answer, ptok, wall = chat( base, mid, msgs, max_tokens=96 ) - dm = tick(K["matched"]) - m0 + dm = reuse_tick() - m0 turn_dms.append(dm) turns_ok &= st == 200 and len(turn_text) > 0 turns_clean &= content_clean(last_answer) @@ -878,7 +974,7 @@ def history_safe(reply: str) -> str: for uq, at in history: msgs.append({"role": "assistant", "content": at}) msgs.append({"role": "user", "content": uq}) - m0 = tick(K["matched"]) + m0 = reuse_tick() st, vtext, vcontent, ptok, wall = chat( base, mid, @@ -886,7 +982,7 @@ def history_safe(reply: str) -> str: max_tokens=96, chat_template_kwargs=template_kwargs, ) - dm = tick(K["matched"]) - m0 + dm = reuse_tick() - m0 rep.check( "variant.status", st == 200 and len(vtext) > 0, @@ -896,12 +992,34 @@ def history_safe(reply: str) -> str: "variant.content_clean", content_clean(vcontent), repr(vcontent[:70]) ) if ck: - rep.check( - "variant.adopted", - dm >= div_floor, - f"matched +{dm} >= grid floor {div_floor} under " - "changed render kwargs", - ) + if dm >= div_floor: + rep.check( + "variant.adopted", + True, + f"matched +{dm} >= grid floor {div_floor} under " + "changed render kwargs", + ) + else: + # Kwargs that rewrite the prompt head (gpt-oss + # reasoning_effort edits the system block) leave no + # shared prefix; the identical resend then proves + # adoption under the new render shape instead. + m0 = reuse_tick() + st2, vtext2, _, _, wall2 = chat( + base, + mid, + msgs, + max_tokens=96, + chat_template_kwargs=template_kwargs, + ) + dm2 = reuse_tick() - m0 + rep.check( + "variant.adopted", + st2 == 200 and dm2 >= div_floor, + f"kwargs change matched +{dm} (head divergence); " + f"resend adopted +{dm2} >= {div_floor}, " + f"{wall2:.1f}s", + ) else: rep.note("variant.matched", f"+{dm}") @@ -1174,6 +1292,207 @@ def fire_short(): f"replay matched +{dm}", ) + # -- sibling burst ------------------------------------------------- + # Unit 4 fresh gate: K siblings sharing a cold user-turn prefix + # arrive together. On the block tier the followers must end warm + # (held for the leader's post-prefill stores). The exact tier + # stores user-turn prefixes only at retirement, beyond the hold + # ceiling, so its numbers are reported, not gated. Ckpt-tier + # models admit one row at a time; no co-admission window exists. + if a.tier == "ckpt": + rep.note( + "burst.scope", + "ckpt tier admits B=1; no co-admission window to gate", + ) + else: + burst_text, n_burst = deep_prefix( + a.prefix_words // 2, header="Sibling-burst variant log.\n" + ) + eb = n_burst // 2 + fb = entry_facts(eb) + plug_text, n_plug = deep_prefix(2000, header="Burst plug log.\n") + m0 = tick(K["matched"]) + fr0 = ((Client(base).metrics()[1] or {}).get("server") + or {}).get("freshness") or {} + + def fire_plug(): + # unrelated prefill in flight while the siblings land, so + # they queue up and co-admit in one formation tick -- the + # exact window the fresh gate exists for + return chat( + base, + mid, + mk_msgs(plug_text + "Say ok."), + max_tokens=16, + ) + + def fire_sib(i): + return chat( + base, + mid, + mk_msgs( + burst_text + + f"Sibling {i}: how many units did entry {eb} " + "report? Answer with just the number." + ), + max_tokens=64, + ) + + t0 = time.monotonic() + with ThreadPoolExecutor(max_workers=a.concurrency + 1) as ex: + fut_plug = ex.submit(fire_plug) + time.sleep(0.25) + futs = [ex.submit(fire_sib, i) for i in range(a.concurrency)] + bres = [f.result() for f in futs] + plug_res = fut_plug.result() + b_wall = time.monotonic() - t0 + settle(K["stores"]) + dm = tick(K["matched"]) - m0 + fr1 = ((Client(base).metrics()[1] or {}).get("server") + or {}).get("freshness") or {} + d_holds = int(fr1.get("holds", 0) or 0) - int( + fr0.get("holds", 0) or 0 + ) + rep.check( + "burst.status", + all(st == 200 and text for st, text, _, _, _ in bres) + and plug_res[0] == 200, + f"{a.concurrency} cold siblings behind a plug prefill, " + f"{b_wall:.1f}s wall", + ) + rep.check( + "burst.coherent", + all( + re.search(rf"\b{fb['reading']}\b", text) + for _, text, _, _, _ in bres + ), + f"entry {eb} -> {fb['reading']} units in every reply", + ) + floor = int(0.6 * (a.prefix_words // 2)) * (a.concurrency - 1) + if a.tier == "block": + rep.check( + "burst.gate_held", + d_holds >= 1, + f"fresh holds +{d_holds} (>=1: the siblings " + "co-admitted and a follower was deferred)", + ) + rep.check( + "burst.followers_warm", + dm >= floor, + f"matched +{dm} >= floor {floor} " + f"({a.concurrency - 1} followers), fresh holds " + f"+{d_holds}", + ) + else: + rep.note( + "burst.followers", + f"matched +{dm} (floor would be {floor}), fresh " + f"holds +{d_holds}; exact tier stores user-turn " + "prefixes at retirement, past the hold ceiling", + ) + rep.note( + "burst.fresh_holds", + f"+{d_holds} (last reason " + f"{str(fr1.get('last_hold_reason'))[:80]!r})", + ) + + # -- queue cap ----------------------------------------------------- + # Unit 5's wire contract. The servers run at cap 1, and depth + # counts past the decode slots (completion batch), so the flood + # must exceed the batch; tail arrivals then land while the queue + # is deep and must be rejected at the socket, never mid-stream. + if a.queue_cap: + flood = 44 + q0 = ((Client(base).metrics()[1] or {}).get("server") + or {}).get("queue") or {} + + def fire_q(i): + return raw_chat( + base, + mid, + mk_msgs( + f"Flood {i}: what is 2 plus {i % 7}? " + "Answer with just the number.", + system=SYSTEM_MSG, + ), + max_tokens=24, + ) + + t0 = time.monotonic() + with ThreadPoolExecutor(max_workers=flood) as ex: + qres = list(ex.map(fire_q, range(flood))) + q_wall = time.monotonic() - t0 + n200 = sum(1 for st, _, _ in qres if st == 200) + rejects = [(h, b) for st, h, b in qres if st == 503] + q1 = ((Client(base).metrics()[1] or {}).get("server") + or {}).get("queue") or {} + d_rej = int(q1.get("rejections", 0) or 0) - int( + q0.get("rejections", 0) or 0 + ) + rep.check( + "queuecap.rejected", + len(rejects) >= 1, + f"{len(rejects)} of {flood} drew 503 at cap 1, " + f"{q_wall:.1f}s wall", + ) + rep.check( + "queuecap.no_other_errors", + n200 + len(rejects) == flood, + f"{n200} served + {len(rejects)} rejected == {flood}", + ) + + def retry_ok(h) -> bool: + try: + return 2 <= int(h.get("Retry-After", "")) <= 60 + except (TypeError, ValueError): + return False + + first_retry = ( + rejects[0][0].get("Retry-After") if rejects else None + ) + rep.check( + "queuecap.retry_after", + bool(rejects) and all(retry_ok(h) for h, _ in rejects), + f"every 503 carries Retry-After in [2, 60] " + f"(first: {first_retry}s)", + ) + err0 = (rejects[0][1].get("error") or {}) if rejects else {} + rep.check( + "queuecap.body_contract", + bool(rejects) + and all( + (b.get("error") or {}).get("type") == "server_overloaded" + and int((b.get("error") or {}).get("queue_cap", 0) or 0) + == 1 + and int((b.get("error") or {}).get("queue_depth", 0) or 0) + >= 1 + for _, b in rejects + ), + "type/queue_cap/queue_depth on every 503 " + f"(first depth {err0.get('queue_depth')})", + ) + rep.check( + "queuecap.metrics_counted", + d_rej >= len(rejects), + f"queue.rejections +{d_rej} >= {len(rejects)} observed 503s", + ) + st, text, _, _, _ = chat( + base, + mid, + mk_msgs( + "What is 2 plus 2? Answer with just the number.", + system="You are a helpful assistant.", + ), + max_tokens=8, + ) + # gate recovery on service, not arithmetic: a log-scoped + # system prompt makes some models refuse off-log questions + rep.check( + "queuecap.recovers", + st == 200 and len(text.strip()) > 0, + f"post-drain retry status {st}: {text[:40]!r}", + ) + # -- bitrate-b isolation ------------------------------------------------ # Only meaningful when a second KV width is requested: the namespaces # must not cross-adopt. Skipped entirely on the fp16 acceptance shape. @@ -1275,11 +1594,12 @@ def fire_short(): chat(base, mid, mk_msgs("Say ok."), max_tokens=4) root = ( - [] if a.no_system else [{"role": "system", "content": SESSION_SYSTEM}] + [] if a.no_system + else [{"role": "system", "content": SESSION_SYSTEM + sfx}] ) root = root + [{"role": "user", "content": prefix + q}] st, text, content, ptok, wall = chat( - base, mid, root, max_tokens=SESSION_REPLY_TOKENS + base, mid, root, max_tokens=a.reply_tokens ) settle(K["stores"]) # the session runs under its own system message, so the log @@ -1329,11 +1649,11 @@ def sess_facts(txt: str) -> bool: {"role": "assistant", "content": summary}, {"role": "user", "content": q}, ] - m0 = tick(K["matched"]) + m0 = reuse_tick() st, text, content, ptok_t, wall = chat( - base, mid, msgs, max_tokens=SESSION_REPLY_TOKENS + base, mid, msgs, max_tokens=a.reply_tokens ) - dm = tick(K["matched"]) - m0 + dm = reuse_tick() - m0 turns_ok &= st == 200 and len(text) > 0 clean_ok &= content_clean(content) comp_row = (st, sess_facts(text), dm, wall) @@ -1430,20 +1750,20 @@ def sess_facts(txt: str) -> bool: 0, ((root_ptok - QUESTION_SLACK) // a.block_size) * a.block_size, ) - m0 = tick(K["matched"]) + m0 = reuse_tick() if stream: st, text, content, events, wall = schat( base, mid, msgs, - max_tokens=1024 if is_abort else SESSION_REPLY_TOKENS, + max_tokens=1024 if is_abort else a.reply_tokens, abort_after=24 if is_abort else None, **extra, ) ptok_t = 0 else: st, text, content, ptok_t, wall = chat( - base, mid, msgs, max_tokens=SESSION_REPLY_TOKENS, **extra + base, mid, msgs, max_tokens=a.reply_tokens, **extra ) events = 0 if kind == "tool" and tool_state == "untried": @@ -1456,11 +1776,11 @@ def sess_facts(txt: str) -> bool: hist = history block = [{"role": "user", "content": dump + "\n" + q_t}] msgs = root + hist + block - m0 = tick(K["matched"]) + m0 = reuse_tick() st, text, content, ptok_t, wall = chat( - base, mid, msgs, max_tokens=SESSION_REPLY_TOKENS + base, mid, msgs, max_tokens=a.reply_tokens ) - dm = tick(K["matched"]) - m0 + dm = reuse_tick() - m0 if is_abort: abort_ok = st == 200 and events >= 24 settle(K["stores"]) @@ -1657,9 +1977,11 @@ def sess_facts(txt: str) -> bool: "prefix_tokens": prefix_tok, "session_turns": a.session or None, "session_final_ptok": sess_final_ptok, + "reply_tokens": a.reply_tokens, "speculative": a.speculative, "system_message": not a.no_system, "system_words": a.system_words or None, + "system_suffix": a.system_suffix, "template_kwargs": template_kwargs, "cold_wall_s": round(cold_wall, 1), "warm_wall_s": round(warm_wall, 1), diff --git a/tests/e2e/run_serve_stress_e2e.py b/tests/e2e/run_serve_stress_e2e.py new file mode 100755 index 0000000..87c6381 --- /dev/null +++ b/tests/e2e/run_serve_stress_e2e.py @@ -0,0 +1,380 @@ +#!/usr/bin/env python3 +"""Serve-path crash hunt: seeded concurrent chaos against one live server. + +Boots the target (APC armed, MTP on request) and runs N worker threads +firing a randomized mix of the seams that have produced crashes: client +aborts mid-stream and mid-prefill, tiny budgets at finish seams, exact +warm resends and sibling bursts through APC admission, growing session +chains with compaction-style rewrites, sampler and seed variety, and +per-request chat_template_kwargs flips. Content quality is not scored; +the pass criterion is purely mechanical: every request returns clean +HTTP, every stream parses, the server log stays free of exceptions, and +the process survives. Every request is journaled with its parameters so +a finding is replayable. + +Usage: ./run_serve_stress_e2e.py --model M.gguf --speculative \ + --minutes 20 --clients 6 --out DIR +""" +from __future__ import annotations + +import argparse +import json +import os +import random +import re +import sys +import threading +import time +import urllib.error +import urllib.request + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from run_apc_depth_e2e import ( # noqa: E402 + deep_prefix, + depth_env, + model_id_of, + serve_args, +) +from server_proc import ServerProc # noqa: E402 + +LOG_BAD = re.compile( + r"Traceback|ERROR|CRITICAL|bad_cast|NotImplementedError|Segmentation|" + r"Fatal|panic|assert", re.IGNORECASE) +# Lines the server emits at ERROR level for CLIENT behavior the chaos +# deliberately causes (mid-stream disconnects); not findings. +LOG_BENIGN = re.compile( + r"ClientDisconnect|Broken pipe|Connection reset|ConnectionResetError") + +PREFIX_SIZES = (800, 2000, 4000, 6000) + + +def build_prefixes(): + """Distinct deep prefixes: deep_prefix is deterministic, so salt each + with a site tag line to keep APC chains distinct across sizes.""" + out = [] + for i, words in enumerate(PREFIX_SIZES): + text, n = deep_prefix(words) + out.append((f"Site tag: node-{i}.\n" + text, n)) + return out + + +class Stats: + def __init__(self): + self.lock = threading.Lock() + self.actions = {} + self.findings = [] + self.anomalies = [] + + def count(self, action): + with self.lock: + self.actions[action] = self.actions.get(action, 0) + 1 + + def finding(self, rec): + with self.lock: + self.findings.append(rec) + print(f" FINDING {json.dumps(rec)[:300]}", flush=True) + + def anomaly(self, rec): + with self.lock: + self.anomalies.append(rec) + + +class Journal: + def __init__(self, path): + self.lock = threading.Lock() + self.fh = open(path, "w") + + def write(self, rec): + with self.lock: + self.fh.write(json.dumps(rec) + "\n") + self.fh.flush() + + +def post_chat(base, body, timeout): + req = urllib.request.Request( + f"{base}/v1/chat/completions", + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}, + method="POST", + ) + return urllib.request.urlopen(req, timeout=timeout) + + +def plain_chat(base, body, timeout): + """Non-streamed request. Returns (status, content, err).""" + try: + with post_chat(base, body, timeout) as resp: + payload = json.loads(resp.read().decode("utf-8", "replace")) + msg = payload["choices"][0]["message"] + return resp.status, str(msg.get("content") or ""), None + except urllib.error.HTTPError as e: + return e.code, e.read().decode("utf-8", "replace")[:400], "http" + except Exception as e: # noqa: BLE001 - every failure is data here + return -1, "", f"{type(e).__name__}: {e}" + + +def stream_chat(base, body, timeout, abort_after=None): + """Streamed request. Returns (status, events, err). abort_after=0 + closes right after connect (mid-prefill cancel); N closes after N + SSE events (mid-decode cancel).""" + body = dict(body) + body["stream"] = True + events = 0 + try: + req = urllib.request.Request( + f"{base}/v1/chat/completions", + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json", + "Accept": "text/event-stream"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=timeout) as resp: + if abort_after == 0: + return resp.status, 0, None + for raw in resp: + line = raw.decode("utf-8", "replace").strip() + if not line.startswith("data:"): + continue + payload = line[5:].strip() + if payload == "[DONE]": + break + json.loads(payload) # malformed SSE JSON is a finding + events += 1 + if abort_after is not None and events >= abort_after: + break + return resp.status, events, None + except urllib.error.HTTPError as e: + return e.code, events, "http" + except json.JSONDecodeError as e: + return -2, events, f"sse-json: {e}" + except Exception as e: # noqa: BLE001 + return -1, events, f"{type(e).__name__}: {e}" + + +def worker(wid, args, base, mid, prefixes, stop_at, stats, journal, rng, + warm_pool, warm_lock): + sessions = [] + i = 0 + while time.monotonic() < stop_at: + i += 1 + action = rng.choices( + ["cold", "warm", "sibling", "session", "stream", "abort", + "tiny", "sampled", "kwargs"], + weights=[14, 10, 10, 16, 14, 16, 6, 8, 6])[0] + stats.count(action) + prefix, n = prefixes[rng.randrange(len(prefixes))] + e = rng.randrange(1, n) + q = (f"Question: what reading did entry {e} report and who logged " + "it? Answer with the number and the name.") + body = {"model": mid, "temperature": 0.0, + "max_tokens": rng.choice([64, 128, 256, 512])} + sys_msg = {"role": "system", + "content": "You are the maintenance log assistant."} + stream = False + abort_after = None + + if action == "cold": + body["messages"] = [sys_msg, + {"role": "user", "content": prefix + q}] + with warm_lock: + warm_pool.append(body["messages"]) + del warm_pool[:-8] + elif action == "warm": + with warm_lock: + msgs = rng.choice(warm_pool) if warm_pool else None + if msgs is None: + continue + body["messages"] = msgs + elif action == "sibling": + tag = rng.randrange(1000) + body["messages"] = [ + sys_msg, + {"role": "user", + "content": prefix + f"Sibling {tag}. " + q}] + elif action == "session": + if sessions and (len(sessions) >= 3 or rng.random() < 0.7): + s = rng.choice(sessions) + if len(s) > 12 and rng.random() < 0.3: + # compaction-style rewrite: drop the middle + del s[2:len(s) - 2] + s.insert(2, {"role": "assistant", + "content": "Summary: readings nominal."}) + s.append({"role": "user", "content": q}) + body["messages"] = list(s) + else: + s = [sys_msg, {"role": "user", "content": prefix + q}] + sessions.append(s) + body["messages"] = list(s) + sessions[:] = sessions[-4:] + elif action == "stream": + stream = True + body["messages"] = [sys_msg, + {"role": "user", "content": prefix + q}] + elif action == "abort": + stream = True + abort_after = rng.choice([0, 1, 3, 8, 20, 40]) + body["messages"] = [sys_msg, + {"role": "user", "content": prefix + q}] + body["max_tokens"] = 1024 + elif action == "tiny": + body["max_tokens"] = rng.randrange(1, 5) + body["messages"] = [sys_msg, + {"role": "user", "content": prefix + q}] + elif action == "sampled": + body.update(temperature=round(rng.uniform(0.5, 1.2), 2), + top_p=round(rng.uniform(0.8, 1.0), 2), + seed=rng.randrange(1 << 30)) + if rng.random() < 0.5: + body["top_k"] = rng.choice([20, 40, 80]) + if rng.random() < 0.3: + body["min_p"] = 0.05 + body["messages"] = [sys_msg, + {"role": "user", "content": prefix + q}] + elif action == "kwargs": + body["chat_template_kwargs"] = { + "enable_thinking": rng.random() < 0.5} + body["messages"] = [sys_msg, + {"role": "user", "content": prefix + q}] + + t0 = time.monotonic() + if stream: + st, events, err = stream_chat(base, body, args.request_timeout, + abort_after=abort_after) + wall = time.monotonic() - t0 + rec = {"wid": wid, "i": i, "action": action, "status": st, + "t": round(time.time(), 1), + "events": events, "wall": round(wall, 1), "err": err, + "abort_after": abort_after, + "max_tokens": body["max_tokens"]} + else: + st, content, err = plain_chat(base, body, args.request_timeout) + wall = time.monotonic() - t0 + rec = {"wid": wid, "i": i, "action": action, "status": st, + "t": round(time.time(), 1), + "wall": round(wall, 1), "err": err, + "max_tokens": body["max_tokens"], + "n_msgs": len(body["messages"])} + if st == 200 and not content and action != "tiny": + stats.anomaly({**rec, "note": "empty content"}) + journal.write(rec) + # An aborted stream reporting a transport error is the abort + # itself; anything else non-200 is a finding. + if action == "abort" and rec["err"] and rec["status"] in (-1, 200): + continue + if rec["status"] != 200 or rec["err"]: + stats.finding(rec) + + +def scan_log(path, stats, seen): + """Scan the server log for new exception lines; benign client-side + disconnect noise is excluded.""" + try: + with open(path, encoding="utf-8", errors="replace") as f: + lines = f.readlines() + except OSError: + return + for ln, line in enumerate(lines): + if ln in seen: + continue + if LOG_BAD.search(line) and not LOG_BENIGN.search(line): + seen.add(ln) + stats.finding({"server_log_line": ln + 1, + "text": line.strip()[:300]}) + elif LOG_BAD.search(line): + seen.add(ln) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--model", required=True) + ap.add_argument("--speculative", action="store_true") + ap.add_argument("--draft-gguf", default=None) + ap.add_argument("--minutes", type=float, default=20.0) + ap.add_argument("--clients", type=int, default=6) + ap.add_argument("--seed", type=int, default=1) + ap.add_argument("--request-timeout", type=float, default=420.0) + ap.add_argument("--out", required=True) + ap.add_argument("--python", default=sys.executable) + args = ap.parse_args() + + os.makedirs(args.out, exist_ok=True) + label = os.path.basename(args.out.rstrip("/")) + log_path = os.path.join(args.out, "server.log") + env = depth_env(os.path.join(args.out, "apc-disk"), 16, 2048, 16) + + print(f"== {label}: building prefixes ==", flush=True) + prefixes = build_prefixes() + stats = Stats() + journal = Journal(os.path.join(args.out, "journal.jsonl")) + seen_log = set() + + with ServerProc( + serve_args(args.model, args.speculative, args.draft_gguf), + env_extra=env, + log_path=log_path, + python=args.python, + ) as srv: + srv.wait_ready(timeout=900) + base = srv.base_url + mid = model_id_of(base, srv) + print(f"== {label}: chaos start, {args.clients} clients x " + f"{args.minutes:.0f} min, seed {args.seed} ==", flush=True) + stop_at = time.monotonic() + args.minutes * 60 + warm_pool, warm_lock = [], threading.Lock() + threads = [ + threading.Thread( + target=worker, + args=(w, args, base, mid, prefixes, stop_at, stats, journal, + random.Random(args.seed * 1000 + w), warm_pool, + warm_lock), + daemon=True) + for w in range(args.clients) + ] + for t in threads: + t.start() + while any(t.is_alive() for t in threads): + time.sleep(10) + scan_log(log_path, stats, seen_log) + if srv.proc.poll() is not None: + stats.finding({"server_died": srv.proc.returncode}) + break + for t in threads: + t.join(timeout=args.request_timeout + 30) + # settle, then final probe + log sweep + time.sleep(3) + st, content, err = plain_chat( + base, {"model": mid, "max_tokens": 8, "temperature": 0.0, + "messages": [{"role": "user", "content": "Say ok."}]}, + 60) + if st != 200: + stats.finding({"post_chaos_probe": st, "err": err}) + scan_log(log_path, stats, seen_log) + if srv.proc.poll() is not None: + stats.finding({"server_died": srv.proc.returncode}) + + total = sum(stats.actions.values()) + report = { + "label": label, + "model": args.model, + "speculative": args.speculative, + "minutes": args.minutes, + "clients": args.clients, + "seed": args.seed, + "requests": total, + "actions": stats.actions, + "findings": stats.findings, + "anomalies": stats.anomalies[:50], + "n_anomalies": len(stats.anomalies), + } + with open(os.path.join(args.out, "report.json"), "w") as f: + json.dump(report, f, indent=1) + print(f" requests {total} actions {json.dumps(stats.actions)}") + print(f" anomalies {len(stats.anomalies)} (empty-content notes)") + verdict = "PASS" if not stats.findings else \ + f"FAIL ({len(stats.findings)} findings)" + print(f"== {label}: {verdict} -> {args.out}/report.json ==", flush=True) + return 0 if not stats.findings else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_apc_pooling.py b/tests/test_apc_pooling.py index a5d73c7..021d929 100644 --- a/tests/test_apc_pooling.py +++ b/tests/test_apc_pooling.py @@ -184,6 +184,46 @@ def __init__(self, model, processor, **kwargs): assert seen["apc_manager"] is mgr +def test_kv_bits_apc_optout_warns_at_boot(monkeypatch, caplog): + """Upstream nulls the APC manager whenever kv_bits is set, with no + signal; a server booted with APC_ENABLED and a KV-quant scheme runs + every request cold. The stash wrapper must say so once at + construction. Warm boots (no kv_bits) and draft-model batches + (upstream nulls their manager by design) stay quiet.""" + import importlib + import logging + from types import SimpleNamespace + + import gmlx.spec_engine as spec_engine + + ar = importlib.import_module("mlx_vlm.generate.ar") + + class _UpstreamLikeBG: + def __init__(self, model, processor, **kwargs): + mgr = kwargs.get("apc_manager") + if mgr is not None and kwargs.get("kv_bits") is not None: + mgr = None + self.apc_manager = mgr + + monkeypatch.setattr(ar, "BatchGenerator", _UpstreamLikeBG) + monkeypatch.setattr(spec_engine, "_SPEC_APC_DISABLED", False) + spec_engine._install_apc_manager_stash() + mgr = object() + + with caplog.at_level(logging.WARNING, logger="gmlx.spec_engine"): + ar.BatchGenerator(SimpleNamespace(), None, apc_manager=mgr, kv_bits=8) + assert any("APC OFF: KV quantization" in r.message for r in caplog.records) + + caplog.clear() + with caplog.at_level(logging.WARNING, logger="gmlx.spec_engine"): + ar.BatchGenerator(SimpleNamespace(), None, apc_manager=mgr) + ar.BatchGenerator(SimpleNamespace(), None, apc_manager=mgr, + kv_bits=8, draft_model=object()) + # the dead-stack probe may warn separately on these bare fakes; only + # the kv_bits line must stay quiet here + assert not any("KV quantization" in r.message for r in caplog.records) + + def test_rebind_to_runtime_origin_recurses_and_skips_ours(): from mlx_lm.models.cache import CacheList, KVCache, RotatingKVCache diff --git a/tests/test_ckpt_decode_lcp.py b/tests/test_ckpt_decode_lcp.py index 18222f4..68c2328 100644 --- a/tests/test_ckpt_decode_lcp.py +++ b/tests/test_ckpt_decode_lcp.py @@ -144,9 +144,11 @@ def test_retirement_uses_newest_snap_at_or_below_lcp(): def test_retirement_rotating_uses_grid_snap(): - """Below the window an off-grid retirement still needs the aligned - decode snapshot (beyond it the full sequence stores directly -- see - test_retirement_rotating_off_grid_beyond_window_stores).""" + """When the whole-sequence store cannot run (replayable prefix + below the full length), the aligned decode snapshot still serves a + rotating retirement. With no cap the same shape now stores its + grid prefix directly -- see + test_retirement_rotating_short_prompt_stores_grid_prefix.""" from test_ckpt_tier import assert_swa_warm_matches, make_swa_cache man = APCManager(num_blocks=64, block_size=16) @@ -155,7 +157,7 @@ def test_retirement_rotating_uses_grid_snap(): cache = make_swa_cache(n, seed=21) snap_src = make_swa_cache(16, seed=22) states = [c for c in snap_src if hasattr(c, "max_size")] - assert retirement_store(man, "ckpt", ids, cache, + assert retirement_store(man, "ckpt", ids, cache, max_len=20, decode_snaps=[(16, states)]) warm, got = ckpt_lookup(man, ids[:16] + [1], extra_hash=0) assert got == 16 and warm is not None diff --git a/tests/test_ckpt_tier.py b/tests/test_ckpt_tier.py index a0348be..6a22785 100644 --- a/tests/test_ckpt_tier.py +++ b/tests/test_ckpt_tier.py @@ -323,7 +323,8 @@ def test_swa_store_lookup_roundtrip(p): def test_swa_store_declines_off_grid_below_window(): """Below the wrap there is no rot tail mechanism: an off-grid store would need a partial window block. Beyond it (see the roundtrip - params) off-grid p stores.""" + params) off-grid p stores. Without grid_truncate the store declines + whole.""" from gmlx.cache_snapshot import _ckpt_stats man = APCManager(num_blocks=64, block_size=16) p = 20 # < W=32 and 20 % 16 != 0 @@ -333,17 +334,104 @@ def test_swa_store_declines_off_grid_below_window(): assert _ckpt_stats(man)["ckpt_declines"] == {"grid": 1} -def test_retirement_rotating_without_snap_declines(): - """No exact-tier fallback: a below-window off-grid rotating - retirement with no decode snapshot stores nothing, and the exact - tier stays empty so the stock warm path never bypasses ckpt - arming.""" +def assert_grid_warm_matches(warm, orig, p): + """warm at truncated p vs the deeper original: rot canonical is the + temporal prefix [0..p), plain KV the same slice.""" + from gmlx.cache_snapshot import rotating_canonical_window + assert len(warm) == len(orig) + for w, o in zip(warm, orig): + if isinstance(o, RotatingKVCache): + kw, vw, mw = rotating_canonical_window(w) + assert mw[2] == p and mw[3] == p + assert mx.array_equal(kw, o.keys[..., :p, :]).item() + assert mx.array_equal(vw, o.values[..., :p, :]).item() + else: + assert int(w.offset) == p + assert mx.array_equal( + w.keys[..., :p, :], o.keys[..., :p, :]).item() + assert mx.array_equal( + w.values[..., :p, :], o.values[..., :p, :]).item() + + +def test_grid_truncate_store_lookup_roundtrip(): + """grid_truncate turns the below-window decline into a terminal + store at b_full; the record is a faithful shorter run.""" + import gmlx.cache_snapshot as cs + man = APCManager(num_blocks=64, block_size=16) + p = 20 + cache = make_swa_cache(p, seed=5) + ids = list(range(300, 300 + p)) + assert ckpt_store(man, ids, cache, extra_hash=2, + grid_truncate=True) == 16 + st = cs._ckpt_stats(man) + assert st["ckpt_grid_truncate"] == 1 + assert st["ckpt_declines"] == {} + rec = next(iter(cs._ckpt_records(man).values())) + assert rec.p == 16 and rec.b_full == 16 + assert rec.ids == tuple(ids[:16]) + warm, got = ckpt_lookup(man, ids + [1], extra_hash=2) + assert got == 16 + assert_grid_warm_matches(warm, cache, 16) + + +def test_grid_truncate_sub_block_declines(): + """b_full < 2: nothing block-aligned to keep, decline as before.""" + from gmlx.cache_snapshot import _ckpt_stats + man = APCManager(num_blocks=64, block_size=16) + p = 10 # b_full = 0 + cache = make_swa_cache(p, seed=5) + assert ckpt_store(man, list(range(300, 300 + p)), cache, + extra_hash=2, grid_truncate=True) == 0 + assert _ckpt_stats(man)["ckpt_declines"] == {"grid": 1} + + +def test_grid_truncate_recurrent_layout_declines(): + """State cannot rewind: an arr layer in the layout keeps the + decline even with grid_truncate.""" + from gmlx.cache_snapshot import _ckpt_stats + man = APCManager(num_blocks=64, block_size=16) + p = 20 + cache = make_swa_cache(p, seed=5) + arr = ArraysCache(size=2) + arr.cache = [mx.random.normal((1, 3, D)), + mx.random.normal((1, H, D, D))] + cache.append(arr) + assert ckpt_store(man, list(range(300, 300 + p)), cache, + extra_hash=2, grid_truncate=True) == 0 + assert _ckpt_stats(man)["ckpt_declines"] == {"grid": 1} + assert all(b.ref_cnt == 0 for b in man.pool) + + +def test_grid_truncate_beyond_window_stores_full(): + """At or beyond the wrap the off-grid store already works whole; + grid_truncate must not truncate it.""" + from gmlx.cache_snapshot import _ckpt_stats + man = APCManager(num_blocks=64, block_size=16) + p = 40 # >= W=32, unaligned + cache = make_swa_cache(p, seed=5) + ids = list(range(300, 300 + p)) + assert ckpt_store(man, ids, cache, extra_hash=2, + grid_truncate=True) == p + assert _ckpt_stats(man)["ckpt_grid_truncate"] == 0 + warm, got = ckpt_lookup(man, ids + [1], extra_hash=2) + assert got == p + assert_swa_warm_matches(warm, cache, p) + + +def test_retirement_rotating_short_prompt_stores_grid_prefix(): + """A below-window off-grid rotating retirement with no decode + snapshot stores the block-grid prefix (grid_truncate), not nothing. + No exact-tier fallback: the exact tier stays empty so the stock + warm path never bypasses ckpt arming.""" man = APCManager(num_blocks=64, block_size=16) p = 20 # < W and unaligned cache = make_swa_cache(p, seed=3) ids = list(range(300, 300 + p)) - assert not retirement_store(man, "ckpt", ids, cache, row=0, - extra_hash=1) + assert retirement_store(man, "ckpt", ids, cache, row=0, + extra_hash=1) == 16 + warm, got = ckpt_lookup(man, ids + [1], extra_hash=1) + assert got == 16 + assert_grid_warm_matches(warm, cache, 16) entry, plen = man.lookup_exact_cache(ids + [1], extra_hash=1) assert entry is None and plen == 0 assert man.stats_snapshot()["exact_stores"] == 0 diff --git a/tests/test_decode_batch.py b/tests/test_decode_batch.py new file mode 100644 index 0000000..b04ac4b --- /dev/null +++ b/tests/test_decode_batch.py @@ -0,0 +1,58 @@ +"""Decode concurrency control: env resolution + serve-path injection.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("mlx_vlm") + +import gmlx.decode_batch as db # noqa: E402 + + +def test_default(monkeypatch): + monkeypatch.delenv("GMLX_DECODE_BATCH", raising=False) + assert db.decode_batch() == db.DEFAULT_DECODE_BATCH + + +def test_env_override(monkeypatch): + monkeypatch.setenv("GMLX_DECODE_BATCH", "4") + assert db.decode_batch() == 4 + + +def test_zero_restores_upstream(monkeypatch): + from mlx_vlm.generate.ar import DEFAULT_COMPLETION_BATCH_SIZE + + monkeypatch.setenv("GMLX_DECODE_BATCH", "0") + assert db.decode_batch() == int(DEFAULT_COMPLETION_BATCH_SIZE) + + +def test_garbage_falls_back(monkeypatch): + monkeypatch.setenv("GMLX_DECODE_BATCH", "lots") + assert db.decode_batch() == db.DEFAULT_DECODE_BATCH + monkeypatch.setenv("GMLX_DECODE_BATCH", "-3") + assert db.decode_batch() == db.DEFAULT_DECODE_BATCH + + +def test_stash_wrapper_injects(monkeypatch): + from mlx_vlm.generate import ar + + import gmlx.spec_engine as spec_engine + + seen = {} + + class _BG: + def __init__(self, model, processor, **kwargs): + seen.update(kwargs) + + monkeypatch.setattr(ar, "BatchGenerator", _BG) + spec_engine._install_apc_manager_stash() + + monkeypatch.setenv("GMLX_DECODE_BATCH", "5") + ar.BatchGenerator(SimpleNamespace(), None) + assert seen["completion_batch_size"] == 5 + + seen.clear() + ar.BatchGenerator(SimpleNamespace(), None, completion_batch_size=3) + assert seen["completion_batch_size"] == 3 diff --git a/tests/test_deepseek_v4_mtp.py b/tests/test_deepseek_v4_mtp.py index 8302978..c2bdce2 100644 --- a/tests/test_deepseek_v4_mtp.py +++ b/tests/test_deepseek_v4_mtp.py @@ -589,3 +589,120 @@ def test_cacheless_forward_matches_cached(): got = lm(prompt).logits assert got.shape == ref.shape assert mx.abs(got.astype(mx.float32) - ref.astype(mx.float32)).max().item() < 1e-5 + + +def test_batch_pooling_cache_verify_sized_undo_armed(): + # Regression: the batch cache stashed its undo only for L <= 2, so a + # block-4 verify write (L=4) cleared it and the first rejection past + # the remainder buffer crashed rollback on every APC-restored (batch + # lifted) request. Parity with the scalar cache: L <= 6. + from gmlx.deepseek_v4_cache import BatchPoolingCache + + mx.random.seed(19) + b = BatchPoolingCache(4, [0, 0]) + b.accumulate_windows( + mx.random.normal((2, 1, 8)), mx.random.normal((2, 1, 2)), 0) + b.accumulate_windows( + mx.random.normal((2, 4, 8)), mx.random.normal((2, 4, 2)), 0) + assert b._can_trim(2) + assert b._can_trim(3) + assert b.trim(2) == 2 + + +def test_batch_pooling_cache_trim_window_recompletion(): + # Batch twin of test_pooling_cache_trim_window_recompletion, including + # the per-row watermark move (pooled rows kept above the restored + # watermark) and a ragged case where only one row re-completes. + from gmlx.deepseek_v4_cache import BatchPoolingCache + + mx.random.seed(23) + ratio, B, D, G = 4, 2, 8, 2 + + def feed(cache, kv, gate, raws, valid=None): + L = kv.shape[1] + if valid is not None: + cache.prepare(lengths=valid) + r_kv, _, _ = cache.accumulate_windows(kv, gate, 0) + if valid is not None: + cache.finalize() + for i in range(B): + vl = L if valid is None else valid[i] + for j in range(vl): + raws[i].append(kv[i, j]) + # Deterministic compressor stand-in fed from the raw history, so + # control and test pooled rows are comparable bit-for-bit. + new_counts = [ + len(raws[i]) // ratio - cache._pool_lengths[i] for i in range(B)] + max_new = max(new_counts) + if max_new > 0: + px = mx.zeros((B, max_new, D)) + for i in range(B): + pl = cache._pool_lengths[i] + for j in range(new_counts[i]): + w = mx.stack( + raws[i][(pl + j) * ratio:(pl + j + 1) * ratio]) + px[i, j] = w.mean(axis=0) + cache.update_and_fetch(px) + + def check(control, test_cache): + assert control.remainder == test_cache.remainder + assert control._pool_lengths == test_cache._pool_lengths + for i in range(B): + r = control.remainder[i] + if r > 0: + assert mx.array_equal( + control.buf_kv[i, :r], test_cache.buf_kv[i, :r]) + pl = control._pool_lengths[i] + if pl > 0: + assert mx.array_equal( + control.pooled[i, :pl], test_cache.pooled[i, :pl]) + if control._prev_kv is not None: + assert test_cache._prev_kv is not None + for i in range(B): + if control._pool_lengths[i] > 0: + assert mx.array_equal( + control._prev_kv[i], test_cache._prev_kv[i]) + + def run(pre_count, upd_len, n_trim, ragged=False): + pre = [(mx.random.normal((B, 1, D)), mx.random.normal((B, 1, G))) + for _ in range(pre_count)] + rag = (mx.random.normal((B, 2, D)), mx.random.normal((B, 2, G))) + upd = mx.random.normal((B, upd_len, D)) + upd_g = mx.random.normal((B, upd_len, G)) + k = upd_len - n_trim + + test_cache = BatchPoolingCache(ratio, [0] * B) + raws = [[] for _ in range(B)] + for kv, g in pre: + feed(test_cache, kv, g, raws) + if ragged: + feed(test_cache, rag[0], rag[1], raws, valid=[2, 1]) + feed(test_cache, upd, upd_g, raws) + assert test_cache._can_trim(n_trim) + assert test_cache.trim(n_trim) == n_trim + + control = BatchPoolingCache(ratio, [0] * B) + raws_c = [[] for _ in range(B)] + for kv, g in pre: + feed(control, kv, g, raws_c) + if ragged: + feed(control, rag[0], rag[1], raws_c, valid=[2, 1]) + if k > 0: + feed(control, upd[:, :k], upd_g[:, :k], raws_c) + check(control, test_cache) + + # rem 2 + 3-wide: trim(1) re-completes (total 4), trim(2) stays in + # the buffer (total 3). + run(2, 3, 1) + run(2, 3, 2) + # rem 3 + 3-wide: both trims re-complete a window. + run(3, 3, 1) + run(3, 3, 2) + # block-4 verify shapes (the live crash): rem 1 + 4-wide. + run(1, 4, 2) + run(3, 4, 2) + run(3, 4, 3) + # ragged remainders [r, r-1]: one row re-completes, the other stays + # in its buffer -- the per-row branch split. + run(2, 3, 1, ragged=True) + run(3, 4, 2, ragged=True) diff --git a/tests/test_fresh_gate.py b/tests/test_fresh_gate.py new file mode 100644 index 0000000..a8f0760 --- /dev/null +++ b/tests/test_fresh_gate.py @@ -0,0 +1,246 @@ +"""Freshness admission gate: hold rules, coverage peek, stash merge.""" + +import threading +from types import SimpleNamespace + +import pytest + +from mlx_vlm.generate import ar + +import gmlx.fresh_gate as fg + + +class FakeBlock: + def __init__(self, token_ids): + self.token_ids = tuple(token_ids) + + +class FakeManager: + block_size = 16 + + def __init__(self): + self.lock = threading.RLock() + self.hash_table = {} + self._exact_cache = {} + self._kq_anchor_cache = {} + + def add_block_chain(self, ids, extra_hash=0): + from mlx_vlm import apc as _apc + parent = _apc.SEED_PARENT_HASH + bs = self.block_size + for i in range(len(ids) // bs): + chunk = tuple(ids[i * bs:(i + 1) * bs]) + h = _apc._hash_tokens(parent, chunk, extra_hash) + self.hash_table[h] = FakeBlock(chunk) + parent = h + + def add_exact(self, ids, extra_hash=0): + self._exact_cache[len(self._exact_cache)] = SimpleNamespace( + token_ids=tuple(ids), extra_hash=extra_hash) + + def add_anchor(self, ids, extra_hash=0): + self._kq_anchor_cache[(tuple(ids), extra_hash)] = (None, 0) + + +def _seq(uid, ids): + return (uid, list(ids), 200, {}, None, None) + + +class FakeGen: + completion_batch_size = 32 + prefill_batch_size = 8 + + def __init__(self, sequences, manager=None): + self.apc_manager = manager if manager is not None else FakeManager() + self._generation_batch = [] + self._prompt_batch = None + self._unprocessed_sequences = list(sequences) + self.admitted = [] + + def _apc_extra_hash(self, prompt_kwargs): + return 0 + + +SHARED = list(range(1000, 1512)) # 512 shared tokens +A = _seq("a", SHARED + [1, 2, 3]) +B = _seq("b", SHARED + [7, 8, 9]) +C = _seq("c", SHARED + [4, 5, 6]) +LONER = _seq("z", list(range(5000, 5512))) + + +def test_lcp(): + assert fg._lcp([1, 2, 3], [1, 2, 4]) == 2 + assert fg._lcp([], [1]) == 0 + long = list(range(20000)) + div = long[:9000] + [-1] + long[9001:] + assert fg._lcp(long, div) == 9000 + assert fg._lcp(long, long) == 20000 + + +def test_holds_uncovered_sibling(): + g = FakeGen([A, B]) + assert fg._keep_count(g) == 1 + + +def test_head_is_never_held(): + g = FakeGen([A]) + assert fg._keep_count(g) is None + + +def test_truncates_at_first_held_follower(): + g = FakeGen([A, B, C]) + assert fg._keep_count(g) == 1 + + +def test_unrelated_follower_admits(): + g = FakeGen([A, LONER]) + assert fg._keep_count(g) is None + + +def test_unrelated_head_then_siblings_cut_after_leader(): + g = FakeGen([LONER, A, B]) + assert fg._keep_count(g) == 2 + + +def test_block_coverage_admits(): + g = FakeGen([A, B]) + g.apc_manager.add_block_chain(SHARED) + assert fg._keep_count(g) is None + + +def test_exact_coverage_admits(): + g = FakeGen([A, B]) + g.apc_manager.add_exact(SHARED) + assert fg._keep_count(g) is None + + +def test_anchor_coverage_admits(): + g = FakeGen([A, B]) + g.apc_manager.add_anchor(SHARED) + assert fg._keep_count(g) is None + + +def test_coverage_under_other_extra_hash_does_not_count(): + g = FakeGen([A, B]) + g.apc_manager.add_exact(SHARED, extra_hash=99) + assert fg._keep_count(g) == 1 + + +def test_short_shared_prefix_admits(): + short = list(range(100, 164)) + g = FakeGen([_seq("a", short + [1]), _seq("b", short + [2])]) + assert fg._keep_count(g) is None + + +def test_kill_switch(monkeypatch): + monkeypatch.setenv("GMLX_APC_FRESH_WAIT_MS", "0") + g = FakeGen([A, B]) + assert fg._keep_count(g) is None + + +def test_hold_ceiling_admits_cold(monkeypatch, caplog): + clock = [0.0] + monkeypatch.setattr( + fg, "time", SimpleNamespace(perf_counter=lambda: clock[0])) + monkeypatch.setenv("GMLX_APC_FRESH_WAIT_MS", "100") + g = FakeGen([A, B]) + assert fg._keep_count(g) == 1 + clock[0] = 0.2 + with caplog.at_level("WARNING"): + assert fg._keep_count(g) is None + assert any("ceiling" in r.message for r in caplog.records) + + +def test_no_hold_while_prompt_batch_live(): + g = FakeGen([A, B]) + g._prompt_batch = object() + assert fg._keep_count(g) is None + + +def test_no_hold_without_manager(): + g = FakeGen([A, B], manager=False) + g.apc_manager = None + assert fg._keep_count(g) is None + + +def test_held_state_pruned_when_request_leaves(): + g = FakeGen([A, B]) + assert fg._keep_count(g) == 1 + g._unprocessed_sequences = [A] + fg._keep_count(g) + assert "b" not in g._kq_fresh_held + + +def _fake_next(self, **kw): + 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) + fg.install_fresh_admission_gate() + yield ar.BatchGenerator._next + + +def test_install_kill_switch(monkeypatch): + monkeypatch.setenv("GMLX_APC_FRESH_WAIT_MS", "0") + monkeypatch.setattr(ar.BatchGenerator, "_next", _fake_next) + fg.install_fresh_admission_gate() + assert ar.BatchGenerator._next is _fake_next + + +def test_wrapper_serializes_siblings(gated): + g = FakeGen([A, B]) + gated(g) + assert g.admitted == ["a"] + assert [s[0] for s in g._unprocessed_sequences] == ["b"] + g.apc_manager.add_block_chain(SHARED) + gated(g) + assert g.admitted == ["a", "b"] + + +def test_wrapper_restores_tail_ahead_of_arrivals(gated, monkeypatch): + g = FakeGen([A, B, C]) + + def _next_with_arrival(self, **kw): + out = _fake_next(self, **kw) + self._unprocessed_sequences.append(_seq("late", [1, 2])) + return out + + monkeypatch.setattr(ar.BatchGenerator, "_next", _next_with_arrival) + fg.install_fresh_admission_gate() + ar.BatchGenerator._next(g) + assert g.admitted == ["a"] + assert [s[0] for s in g._unprocessed_sequences] == ["b", "c", "late"] + + +def test_wrapper_decision_failure_degrades(gated, monkeypatch): + def _boom(gen): + raise RuntimeError("peek broke") + + monkeypatch.setattr(fg, "_keep_count", _boom) + g = FakeGen([A, B]) + gated(g) + assert g.admitted == ["a", "b"] + + +def test_covered_len_prefers_longest_tier(): + m = FakeManager() + m.add_block_chain(SHARED[:256]) + m.add_exact(SHARED[:300]) + m.add_anchor(SHARED[:400]) + ids = SHARED + [1, 2] + assert fg._covered_len(m, ids, 0) == 400 + + +def test_fresh_stats_counts_holds(): + before = fg.fresh_stats()["holds"] + g = FakeGen([A, B]) + fg._keep_count(g) + st = fg.fresh_stats() + assert st["holds"] == before + 1 + assert "shares" in st["last_hold_reason"] diff --git a/tests/test_mem_preflight.py b/tests/test_mem_preflight.py new file mode 100644 index 0000000..53de026 --- /dev/null +++ b/tests/test_mem_preflight.py @@ -0,0 +1,177 @@ +"""Memory preflight: geometry estimator and the decision table.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("mlx_vlm") + +import gmlx.mem_preflight as mp # noqa: E402 +from mlx_vlm.server.generation import PromptTooLongError # noqa: E402 + + +def _model(**cfg): + return SimpleNamespace(config=SimpleNamespace(**cfg)) + + +DENSE = _model(num_hidden_layers=4, num_attention_heads=8, + num_key_value_heads=8, head_dim=64) +GQA = _model(num_hidden_layers=4, num_attention_heads=8, + num_key_value_heads=2, head_dim=64) +MLA = _model(num_hidden_layers=4, kv_lora_rank=512, qk_rope_head_dim=64) +SWA = _model(num_hidden_layers=6, num_attention_heads=8, + num_key_value_heads=2, head_dim=64, sliding_window=128, + sliding_window_pattern=3) +TYPED = _model(num_hidden_layers=2, num_attention_heads=8, + num_key_value_heads=2, head_dim=64, sliding_window=128, + layer_types=["sliding_attention", "full_attention"]) + + +def test_dense_cost(): + costs = mp.kv_layer_costs(DENSE) + # 2 (K+V) x 8 heads x 64 dim x 2 bytes = 2048 B/token/layer + assert costs == [(None, 2048.0)] * 4 + assert mp.prompt_kv_bytes(costs, 1000) == 4 * 2048.0 * 1000 + + +def test_gqa_cost_scales_with_kv_heads(): + assert mp.kv_layer_costs(GQA) == [(None, 512.0)] * 4 + + +def test_head_dim_derived_from_hidden_size(): + m = _model(num_hidden_layers=1, num_attention_heads=8, + hidden_size=512) + assert mp.kv_layer_costs(m) == [(None, 2 * 8 * 64 * 2.0)] + + +def test_mla_prices_the_latent(): + costs = mp.kv_layer_costs(MLA) + assert costs == [(None, (512 + 64) * 2.0)] * 4 + + +def test_sliding_pattern_caps_layers(): + costs = mp.kv_layer_costs(SWA) + # pattern 3: layers 3 and 6 are global, the rest cap at the window + assert [w for w, _ in costs] == [128, 128, None, 128, 128, None] + per = 512.0 + assert mp.prompt_kv_bytes(costs, 1000) == per * (4 * 128 + 2 * 1000) + + +def test_layer_types_win_over_pattern(): + assert [w for w, _ in mp.kv_layer_costs(TYPED)] == [128, None] + + +def test_kv_bits_lower_the_cost(): + assert mp.kv_layer_costs(DENSE, bytes_per_elem=1.0)[0][1] == 1024.0 + + +def test_unprobeable_geometry_is_none(): + assert mp.kv_layer_costs(SimpleNamespace(config=None)) is None + assert mp.kv_layer_costs(_model(num_hidden_layers=4)) is None + + +def _rg(model, kv_bits=None, tokens=None): + rg = SimpleNamespace(model=model, kv_bits=kv_bits) + calls = [] + + def _pre(prompt, *a, **k): + calls.append(prompt) + return {"input_ids": [0] * (tokens or len(prompt))} + rg._preprocess_request = _pre + rg._pre_calls = calls + return rg + + +@pytest.fixture +def tight(monkeypatch): + """A box where 100k tokens of DENSE KV (0.8 GB) does not fit.""" + monkeypatch.setattr(mp, "available_drained_bytes", lambda: 0.5e9) + monkeypatch.setattr( + "gmlx.prefill_decay.score_transient_bytes", + lambda model, pc, depth: 0.0) + + +def test_prompt_impossible_rejects(tight): + rg = _rg(DENSE, tokens=100_000) + with pytest.raises(PromptTooLongError) as e: + mp.preflight_prompt_memory(rg, "x" * 200_000) + assert "cannot fit" in str(e.value) + assert "prompt_tokens=100000" in str(e.value) + + +def test_big_but_possible_admits(tight): + rg = _rg(DENSE, tokens=50_000) # 0.4 GB < 0.5 GB + mp.preflight_prompt_memory(rg, "x" * 200_000) + assert rg._pre_calls # char bound failed, token count decided + + +def test_small_prompt_never_tokenizes(tight): + rg = _rg(DENSE) + mp.preflight_prompt_memory(rg, "short prompt") + assert not rg._pre_calls + + +def test_pinned_max_impossible_rejects(tight, monkeypatch): + monkeypatch.setattr( + "mlx_vlm.server.generation.get_server_max_tokens", lambda: 512) + rg = _rg(DENSE, tokens=30_000) # 0.24 GB prompt + args = SimpleNamespace(max_tokens=40_000) # +0.32 GB pinned + with pytest.raises(PromptTooLongError) as e: + mp.preflight_prompt_memory(rg, "x" * 60_000, args=args) + assert "max_tokens" in str(e.value) + + +def test_default_max_never_generation_rejects(tight, monkeypatch): + monkeypatch.setattr( + "mlx_vlm.server.generation.get_server_max_tokens", lambda: 40_000) + rg = _rg(DENSE, tokens=30_000) + args = SimpleNamespace(max_tokens=40_000) # equals server default + mp.preflight_prompt_memory(rg, "x" * 60_000, args=args) + + +def test_kill_switch(tight, monkeypatch): + monkeypatch.setenv("GMLX_PREFLIGHT_MEM", "0") + rg = _rg(DENSE, tokens=100_000) + mp.preflight_prompt_memory(rg, "x" * 200_000) + + +def test_media_requests_skip(tight): + rg = _rg(DENSE, tokens=100_000) + mp.preflight_prompt_memory(rg, "x" * 200_000, images=["img"]) + + +def test_no_model_skips(tight): + rg = _rg(None) + mp.preflight_prompt_memory(rg, "x" * 200_000) + + +def test_probe_failure_admits(tight, monkeypatch): + rg = _rg(DENSE) + + def _boom(prompt, *a, **k): + raise RuntimeError("tokenizer broke") + + rg._preprocess_request = _boom + mp.preflight_prompt_memory(rg, "x" * 200_000) + + +def test_kv_bits_shrink_the_estimate(tight): + # 8-bit KV halves the need: 100k tokens fit where fp16 does not + rg = _rg(DENSE, kv_bits=8, tokens=55_000) + mp.preflight_prompt_memory(rg, "x" * 200_000) + + +def test_install_wraps_both_and_is_idempotent(monkeypatch): + from mlx_vlm.server.generation import ResponseGenerator as RG + + saved_gen, saved_val = RG.generate, RG.validate_context_budget + try: + mp.install_memory_preflight() + g1, v1 = RG.generate, RG.validate_context_budget + assert g1 is not saved_gen and v1 is not saved_val + mp.install_memory_preflight() + assert RG.generate is g1 and RG.validate_context_budget is v1 + finally: + RG.generate, RG.validate_context_budget = saved_gen, saved_val diff --git a/tests/test_mtp_preempt_resume.py b/tests/test_mtp_preempt_resume.py index b2d99ea..a4b6f7c 100644 --- a/tests/test_mtp_preempt_resume.py +++ b/tests/test_mtp_preempt_resume.py @@ -174,6 +174,40 @@ def test_armless_capture_skips_entry_seed(): assert len(d.prefill_calls) == 1 +class _WrongDraftDrafter(_ArmableDrafter): + """One deliberately wrong draft token per round, so every armed round + rejects at position 0 and takes the rollback seam.""" + + def draft_block(self, b, hidden, kv, n, sampler, dtype, **kw): + self.draft_calls.append(int(b.shape[0])) + return ((b[:, None] + 5) % VOCAB).astype(dtype) + + +class _ScalarRollbackLM(_VerifyEchoLM): + """int(accepted) is the dsv4/muse scalar-only rollback contract: a + one-element list raises TypeError instead of recording.""" + + def __init__(self): + super().__init__() + self.rollback_accepted = [] + + def rollback_speculative_cache(self, prompt_cache, gdn_states, accepted, + block_size): + self.rollback_accepted.append(int(accepted)) + + +def test_width_one_rejection_calls_scalar_rollback(): + """A width-1 batch spec round that rejects a draft must hand the rollback + hook that row's int, not a one-element list (models defining the hook are + B=1-limited, so wider spec rounds cannot reach it).""" + d = _WrongDraftDrafter(cap=1) + lm = _ScalarRollbackLM() + out, _, _ = _drive_armless(d, B=1, max_tokens=4, lm=lm) + assert lm.rollback_accepted and all(a == 0 for a in lm.rollback_accepted) + # the stream stays the plain echo chain despite every-round rejection + assert [toks for toks, _ in out] == [[2], [3], [4]] + + # -- resume (gated batch drains under the cap) ---------------------------- @@ -181,6 +215,19 @@ def _finish_row0(orig, tok): return orig == 0 +def test_row_finish_filters_without_shared_kv(): + """A drafter with uses_shared_kv=False leaves next_shared_kv None; the + first mid-batch row finish takes the keep-slots filter, which must not + iterate it (the qwen35 owned-MTP concurrent-serve crash).""" + d = _ArmableDrafter(cap=4) + d.uses_shared_kv = False + out, lm, cache = _drive_armless( + d, B=2, max_tokens=40, rounds=6, stop_check=_finish_row0) + assert cache[0].width == 1 # row 0 filtered out, loop kept running + assert all(toks[0] is None for toks, _ in out[1:]) + assert all(toks[1] is not None for toks, _ in out) + + def test_resume_after_drain_rearms_and_streams(): """B=3 over cap=2 gates at formation; row 0 finishing drains the batch to the cap. The next round consumes the dispatched plain lookahead (its KV @@ -480,3 +527,171 @@ def test_preempt_waits_for_first_delivery(monkeypatch): assert model._kq_rebuild_emitted == [1] assert [(r.uid, r.token) for r in responses] == [ (7, 12), (0, 101), (7, 101)] + + +# -- preempt capture: the closed round's undelivered tail ------------------ + + +class _EchoDrafter(_ArmableDrafter): + """Draft the echo target's own continuation, so every draft accepts and + a round yields several tokens (the mid-round preempt window).""" + + def draft_block(self, b, hidden, kv, n, sampler, dtype, **kw): + if not isinstance(b, mx.array): + b = mx.array([b]) + b0 = b.reshape(-1).astype(mx.int32) + return mx.stack([(b0 + 1) % VOCAB, (b0 + 2) % VOCAB], + axis=1).astype(dtype) + + +def test_scalar_close_captures_undelivered_tail(): + """A preempt close mid-round hands the verified undelivered tokens to + model._kq_preempt_capture instead of trimming them; a plain close (no + capture attr) keeps the old trim.""" + d = _EchoDrafter(cap=0) + lm = _VerifyEchoLM() + model = SimpleNamespace() + shared = {"full": (mx.zeros((1, 2, 4, 4)), mx.zeros((1, 2, 4, 4)))} + gen = spec._owned_decode_rounds( + model, d, lm, [_FakeCache(width=1)], + hidden=mx.zeros((1, 4, 8)), b=5, shared_kv=shared, + seed_tokens=None, emitted=1, max_tokens=20, sampler=None, + draft_block_size=None) + assert next(gen) == 6 # round is [6, 7, 8]; one delivered + captured = [] + model._kq_preempt_capture = captured + gen.close() + assert captured == [7, 8] + + +def _capturing_rounds(calls): + """_recording_rounds plus the capture contract: a preempt close hands + two undelivered tokens to model._kq_preempt_capture.""" + + inner = _recording_rounds(calls) + + def fake_rounds(model, draft_model, prompt_cache, hidden, **kw): + it = inner(model, draft_model, prompt_cache, hidden, **kw) + n = 0 + try: + for toks, meta in it: + n += 1 + yield toks, meta + except GeneratorExit: + cap = getattr(model, "_kq_preempt_capture", None) + if cap is not None: + base = 100 * len(calls) + cap.extend([base + n + 1, base + n + 2]) + it.close() + raise + + return fake_rounds + + +def test_preempt_delivers_captured_tail(monkeypatch): + """The captured tail reaches the wire ahead of the waiter's admission, + and the rebuild resumes from the tail's last token (the round's bonus, + whose KV is not in the cache).""" + from mlx_vlm.generate import ar + from gmlx.spec_engine import install_continuous_batch_admission + + install_continuous_batch_admission() + calls = [] + monkeypatch.setattr(ar, "run_speculative_server_rounds", + _capturing_rounds(calls)) + + model = SimpleNamespace() + host = _make_batch(ar, uids=(0,), model=model) + assert [r.token for r in host.next()] == [5] + assert [r.token for r in host.next()] == [101] + + host.extend(_make_batch(ar, uids=(7,), model=model)) + responses = host.next() + assert calls[0]["closed"] is True + assert len(calls) == 2 + assert calls[1]["first_bonus"].tolist() == [103] + # first token + round token + two captured + assert model._kq_rebuild_emitted == [4] + assert [(r.uid, r.token) for r in responses] == [ + (0, 102), (0, 103), (7, 12), (0, 201), (7, 201)] + + +def test_preempt_captured_tail_finishes_row(monkeypatch): + """A captured token that exhausts the budget finishes the row: the tail + truncates at the finish, no rebuild happens, and the waiter promotes.""" + from mlx_vlm.generate import ar + from gmlx.spec_engine import install_continuous_batch_admission + + install_continuous_batch_admission() + calls = [] + monkeypatch.setattr(ar, "run_speculative_server_rounds", + _capturing_rounds(calls)) + + model = SimpleNamespace() + host = _make_batch(ar, uids=(0,), model=model, max_tokens=3) + host.next() + host.next() + host.extend(_make_batch(ar, uids=(7,), model=model)) + + responses = host.next() + assert len(calls) == 1 # no rebuild + assert responses[0].uid == 0 + assert responses[0].token == 102 + assert responses[0].finish_reason == "length" + # 103 dropped past the finish; the waiter promoted and sent its first + assert [(r.uid, r.token) for r in responses[1:]] == [(7, 12)] + + +# -- gated injection: double buffer consumed, never discarded -------------- + + +class _InputLogEchoLM(_VerifyEchoLM): + def __init__(self): + super().__init__() + self.plain_inputs = [] + + def __call__(self, x, cache=None, return_hidden=False, + return_shared_kv=False, **kw): + if not return_hidden: + self.plain_inputs.append(x.reshape(-1).tolist()) + return super().__call__(x, cache=cache, return_hidden=return_hidden, + return_shared_kv=return_shared_kv, **kw) + + +def test_injection_consumes_gated_double_buffer(): + """An injection landing while the gated buffer holds a dispatched step + consumes that step (no forward) and admits at the next boundary. The + old discard re-forwarded the step's inputs, double-writing their KV + and skipping one delivered token per row.""" + d = _StrictDrafter(cap=1) + model = SimpleNamespace() + lm = _InputLogEchoLM() + prompt_cache = [_FakeCache(width=2)] + gen = _owned_decode_rounds_batch( + model, d, lm, prompt_cache, + hidden=None, b=[1, 2], shared_kv=None, seed_tokens=None, + emitted=[1, 1], max_tokens=10, sampler=None, + draft_block_size=None) + toks, _ = next(gen) + assert toks == [2, 3] + assert lm.plain_inputs == [[1, 2], [2, 3]] # prime + dispatch + + model._generator_injections = [{ + "uids": ["w"], + "prompt_cache": [_FakeCache(width=1, offset=9)], + "hidden": mx.zeros((1, 1, 8)), + "prompt_tokens": mx.zeros((1, 4), dtype=mx.int32), + "first_tokens": mx.array([7], dtype=mx.int32), + "first_tokens_list": [7], + "shared_kv_states": None, + }] + toks, _ = next(gen) + # hold round: the buffered step is emitted with no new forward + assert toks == [3, 4] + assert lm.plain_inputs == [[1, 2], [2, 3]] + + toks, _ = next(gen) + gen.close() + # admission landed; every forwarded input advanced exactly once per row + assert toks == [4, 5, 8] + assert lm.plain_inputs[2][:2] == [3, 4] diff --git a/tests/test_mtp_width_cap.py b/tests/test_mtp_width_cap.py index 8303124..dc48043 100644 --- a/tests/test_mtp_width_cap.py +++ b/tests/test_mtp_width_cap.py @@ -406,6 +406,26 @@ def test_b_equals_cap_one_speculates(): _drive(d, B=1, max_tokens=3) +class _B1OnlyDrafter(_StrictDrafter): + """reset refuses any left_padding list, even the trivial [0] -- the + dsv4/muse B=1-only contract.""" + + def reset(self, model, left_padding=None): + if left_padding is not None: + raise NotImplementedError("B=1 only") + return super().reset(model, left_padding=left_padding) + + +def test_b1_only_drafter_arms_bare_at_width_one(): + """A width-1 batch (a preempted scalar rebuilds into the batch loop at + B=1) must arm a B=1-only drafter with a bare reset, not crash on the + trivial [0] padding list.""" + d = _B1OnlyDrafter(cap=1, limit=1) + with pytest.raises(AssertionError, match="draft_block"): + _drive(d, B=1, max_tokens=3) + assert d.reset_calls == [None] + + def test_env_zero_ungates_a_capped_drafter(monkeypatch): monkeypatch.setenv("GMLX_MTP_WIDTH_CAP", "0") d = _StrictDrafter(cap=2) diff --git a/tests/test_muse_glimmer_vision.py b/tests/test_muse_glimmer_vision.py index 6e437f1..94abc0d 100644 --- a/tests/test_muse_glimmer_vision.py +++ b/tests/test_muse_glimmer_vision.py @@ -174,7 +174,9 @@ def masked_reference(pixel_values): got, ref = model(img), masked_reference(img) mx.eval(got, ref) err = float(mx.abs(got - ref).max().item()) - assert err < 1e-4, (grid_h, grid_w, err) + # tf32 gemm noise on M5 reaches ~4e-4 at (7,5); a partition or + # reshape bug is orders larger. + assert err < 1e-3, (grid_h, grid_w, err) # --- pixel shuffle ------------------------------------------------------------ diff --git a/tests/test_prefix_cache_budget.py b/tests/test_prefix_cache_budget.py index 24c5e5f..dbde418 100644 --- a/tests/test_prefix_cache_budget.py +++ b/tests/test_prefix_cache_budget.py @@ -205,3 +205,24 @@ def test_clear_resets_bytes(self): cache.clear() assert cache.total_bytes == 0 assert len(cache) == 0 + + +class TestSpecPrefixStats: + def test_counters_and_reset_contract(self): + import gmlx.prefix_cache as pc + + pc.spec_prefix_stats_clear() + cache = SpecPrefixCache() + ids = mx.array(list(range(40))) + cache.store(ids, [_fill_kv(40)], mx.zeros((1, 1, 4))) + hit = cache.lookup(mx.array(list(range(41)))) + assert hit is not None + s = pc.spec_prefix_stats() + assert s["spec_prefix_stores"] == 1 + assert s["spec_prefix_hits"] == 1 + assert s["spec_prefix_hit_tokens"] == 40 + # /v1/cache/reset drops live entries; counters clear separately + pc.clear_all_spec_prefix_caches() + assert len(cache) == 0 + pc.spec_prefix_stats_clear() + assert pc.spec_prefix_stats()["spec_prefix_hits"] == 0 diff --git a/tests/test_queue_cap.py b/tests/test_queue_cap.py new file mode 100644 index 0000000..fa0b72e --- /dev/null +++ b/tests/test_queue_cap.py @@ -0,0 +1,191 @@ +"""Queue depth cap: decision, Retry-After estimate, route rejection.""" + +from __future__ import annotations + +import importlib +from types import SimpleNamespace + +import pytest + +pytest.importorskip("mlx_vlm") + +import gmlx.queue_cap as qc # noqa: E402 + +_APP = importlib.import_module("mlx_vlm.server.app") +_RUNTIME = importlib.import_module("mlx_vlm.server.runtime").runtime + + +def _metrics(done=0, toks=0, gen_toks=0, decode_s=0.0): + return SimpleNamespace( + _requests_completed=done, + _completion_tokens_total=toks, _generated_tokens_total=gen_toks, + _decode_time_total_s=decode_s) + + +def _rg(qsize=0): + return SimpleNamespace(requests=SimpleNamespace(qsize=lambda: qsize)) + + +def _census(monkeypatch, rg_qsize=0, pending=0): + """Wire a fake waiting census: server queue size + generator pending.""" + monkeypatch.setattr(_RUNTIME, "response_generator", _rg(rg_qsize), + raising=False) + gen = SimpleNamespace(_unprocessed_sequences=[object()] * pending) + monkeypatch.setattr(qc, "_GEN_REF", lambda: gen) + + +def test_default_cap_formula(): + assert qc._cap() == 2 * qc._decode_concurrency() + + +def test_cap_tracks_decode_batch(monkeypatch): + monkeypatch.delenv("GMLX_QUEUE_DEPTH_CAP", raising=False) + monkeypatch.setenv("GMLX_DECODE_BATCH", "6") + assert qc._cap() == 12 + + +def test_cap_env_override(monkeypatch): + monkeypatch.setenv("GMLX_QUEUE_DEPTH_CAP", "7") + assert qc._cap() == 7 + + +def test_depth_counts_queue_and_pending(monkeypatch): + gen = SimpleNamespace(_unprocessed_sequences=[1, 2, 3]) + monkeypatch.setattr(qc, "_GEN_REF", lambda: gen) + assert qc._waiting_depth(_rg(qsize=2)) == 5 + + +def test_depth_survives_dead_generator(monkeypatch): + monkeypatch.setattr(qc, "_GEN_REF", lambda: None) + assert qc._waiting_depth(_rg(qsize=1)) == 1 + + +def test_depth_survives_broken_qsize(monkeypatch): + monkeypatch.setattr(qc, "_GEN_REF", None) + rg = SimpleNamespace( + requests=SimpleNamespace(qsize=lambda: 1 / 0)) + assert qc._waiting_depth(rg) == 0 + + +def test_census_publisher_stashes_generator(): + from mlx_vlm.generate import ar as _ar + + qc._install_census() + assert getattr(_ar.BatchGenerator._next, qc._PUB_FLAG, False) + + class _Stub: # SimpleNamespace refuses weakrefs + _unprocessed_sequences = [1] + + stub = _Stub() + try: + _ar.BatchGenerator._next(stub) + except Exception: + pass # stock body needs real state; publish happens first + assert qc._GEN_REF() is stub + + +def test_retry_after_no_stats_is_static(): + assert qc._retry_after_s(_metrics(), 100) == qc._RETRY_DEFAULT_S + + +def test_retry_after_estimate_and_clamp(): + # 10 waiting x 200 mean tokens / 100 tok/s = 20 s + m = _metrics(done=5, toks=1000, gen_toks=5000, decode_s=50.0) + assert qc._retry_after_s(m, 10) == 20 + assert qc._retry_after_s(m, 1) == qc._RETRY_MIN_S + assert qc._retry_after_s(m, 1000) == qc._RETRY_MAX_S + + +def test_check_below_cap_admits(monkeypatch): + _census(monkeypatch, rg_qsize=0, pending=1) + monkeypatch.setattr(_RUNTIME, "metrics", _metrics(), raising=False) + assert qc.check_queue_depth() is None + + +def test_check_at_cap_rejects(monkeypatch): + monkeypatch.setenv("GMLX_QUEUE_DEPTH_CAP", "4") + _census(monkeypatch, rg_qsize=1, pending=3) + monkeypatch.setattr(_RUNTIME, "metrics", _metrics(), raising=False) + resp = qc.check_queue_depth() + assert resp is not None and resp.status_code == 503 + assert resp.headers["retry-after"] == str(qc._RETRY_DEFAULT_S) + + +def test_check_disabled_by_zero(monkeypatch): + monkeypatch.setenv("GMLX_QUEUE_DEPTH_CAP", "0") + _census(monkeypatch, rg_qsize=999, pending=0) + assert qc.check_queue_depth() is None + + +def test_check_no_engine_admits(monkeypatch): + monkeypatch.setenv("GMLX_QUEUE_DEPTH_CAP", "1") + monkeypatch.setattr(_RUNTIME, "response_generator", None, + raising=False) + assert qc.check_queue_depth() is None + + +@pytest.fixture +def app_routes(): + saved = list(_APP.app.router.routes) + yield _APP.app + _APP.app.router.routes[:] = saved + + +def test_route_rejects_with_503_body(app_routes, monkeypatch): + from fastapi.testclient import TestClient + + monkeypatch.setenv("GMLX_QUEUE_DEPTH_CAP", "2") + _census(monkeypatch, rg_qsize=2, pending=3) + monkeypatch.setattr(_RUNTIME, "metrics", + _metrics(done=2, toks=100, + gen_toks=100, decode_s=10.0), + raising=False) + qc.install_queue_depth_cap() + client = TestClient(_APP.app) + r = client.post("/v1/chat/completions", + json={"model": "m", "messages": [ + {"role": "user", "content": "hi"}]}) + assert r.status_code == 503 + body = r.json()["error"] + assert body["queue_cap"] == 2 and body["queue_depth"] == 5 + assert body["type"] == "server_overloaded" + assert "Retry-After" in r.headers + assert 2 <= int(r.headers["Retry-After"]) <= 60 + assert qc.queue_cap_stats()["rejections"] >= 1 + + +def test_route_at_cap_minus_one_serves(app_routes, monkeypatch): + """A request admitted below the cap reaches the stock handler.""" + from fastapi.testclient import TestClient + + monkeypatch.setenv("GMLX_QUEUE_DEPTH_CAP", "6") + _census(monkeypatch, rg_qsize=5, pending=0) + monkeypatch.setattr(_RUNTIME, "metrics", _metrics(), raising=False) + seen = [] + + async def _stub(*a, **k): + seen.append(1) + return {"ok": True} + + for path in ("/v1/chat/completions",): + route = next(r for r in _APP.app.router.routes + if getattr(r, "path", None) == path + and "POST" in (getattr(r, "methods", None) or ())) + import inspect + _stub.__signature__ = inspect.signature(route.endpoint) + from gmlx.server_patches._common import _remove_routes + _remove_routes(_APP.app, path) + _APP.app.add_api_route(path, _stub, methods=["POST"], + include_in_schema=False) + qc.install_queue_depth_cap() + client = TestClient(_APP.app) + r = client.post("/v1/chat/completions", + json={"model": "m", "messages": []}) + assert r.status_code == 200 and seen == [1] + + +def test_install_idempotent(app_routes): + qc.install_queue_depth_cap() + n = len(_APP.app.router.routes) + qc.install_queue_depth_cap() + assert len(_APP.app.router.routes) == n diff --git a/tests/test_seed_rows.py b/tests/test_seed_rows.py new file mode 100644 index 0000000..3491aba --- /dev/null +++ b/tests/test_seed_rows.py @@ -0,0 +1,192 @@ +"""Per-request seed: keyed rows, byte-identical unseeded rows, plumbing.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import mlx.core as mx +import pytest + +pytest.importorskip("mlx_vlm") + +from mlx_vlm.generate import ar # noqa: E402 +from mlx_vlm.server import generation as gen_mod # noqa: E402 +from mlx_vlm.server.generation import _position_keys # noqa: E402 + +import gmlx.seed_rows as sr # noqa: E402 +import gmlx.speculative as spec # noqa: E402 +from gmlx.server_patches.sampling import _FastPositionedSampler # noqa: E402 + + +def _sampler(**kw): + kw.setdefault("temperature", 1.0) + return _FastPositionedSampler(**kw) + + +def _logprobs(rows=2, vocab=64, seed=0): + x = mx.random.uniform(shape=(rows, vocab), key=mx.random.key(seed)) + return x - mx.logsumexp(x, axis=-1, keepdims=True) + + +def test_unseeded_rows_keep_stock_keys(): + s = _sampler() + s._kq_row_seeds["other"] = 999 # registry active, rows unseeded + s._kq_rows = ["a", "b"] + keys = s._row_keys([0, 0], [5, 6]) + assert mx.array_equal(keys, _position_keys(s.seed, [0, 0], [5, 6])) + + +def test_seeded_row_gets_its_own_key(): + s = _sampler() + s._kq_row_seeds["a"] = 123 + s._kq_rows = ["a", "b"] + keys = s._row_keys([0, 0], [5, 5]) + stock = _position_keys(s.seed, [0, 0], [5, 5]) + assert not mx.array_equal(keys[0], stock[0]) + assert mx.array_equal(keys[1], stock[1]) + assert mx.array_equal( + keys[0], mx.random.key(gen_mod._position_seed(123, 0, 5))) + + +def test_rows_context_mismatch_falls_back(): + s = _sampler() + s._kq_row_seeds["a"] = 123 + s._kq_rows = ["a"] # 1 row context, 2-row draw + keys = s._row_keys([0, 0], [5, 6]) + assert mx.array_equal(keys, _position_keys(s.seed, [0, 0], [5, 6])) + + +def test_solo_replay_is_exact(): + lp = _logprobs(rows=1) + + def run(): + s = _sampler(top_p=0.9) + s._kq_row_seeds[7] = 42 + s._kq_rows = [7] + return [int(s.sample_target(lp, row_ids=[0], positions=[p]).item()) + for p in range(1, 12)] + + assert run() == run() + + +def test_mixed_batch_unseeded_row_matches_solo_stock(): + lp = _logprobs(rows=2) + mixed = _sampler() + mixed._kq_row_seeds[7] = 42 + mixed._kq_rows = [7, 8] + out = mixed.sample_target(lp, row_ids=[0, 0], positions=[3, 3]) + stock = _sampler() + ref = stock.sample_target(lp, row_ids=[0, 0], positions=[3, 3]) + assert int(out[1].item()) == int(ref[1].item()) + + +def test_registry_caps(): + s = _sampler() + for i in range(sr._MAX_SEEDS + 10): + sr.register_row_seed(s, i, i) + assert len(s._kq_row_seeds) == sr._MAX_SEEDS + assert 0 not in s._kq_row_seeds + assert sr._MAX_SEEDS + 9 in s._kq_row_seeds + + +def test_seeded_target_draw_replays(): + lp = _logprobs(rows=4) + s = _sampler() + s._kq_row_seeds["u"] = 42 + s._kq_rows = ["u"] + a = spec._seeded_target_draw(s, lp, base_pos=10) + b = spec._seeded_target_draw(s, lp, base_pos=10) + assert mx.array_equal(a, b) + assert s._kq_rows == ["u"] # context restored + + +def test_seeded_target_draw_unseeded_uses_process_stream(): + lp = _logprobs(rows=4) + calls = [] + orig_call = _FastPositionedSampler.__call__ + + class Probe(_FastPositionedSampler): + def __call__(self, logprobs): + calls.append(1) + return orig_call(self, logprobs) + + p = Probe(temperature=1.0) + p._kq_rows = ["u"] # no seed registered + spec._seeded_target_draw(p, lp, base_pos=10) + assert calls == [1] + + +@pytest.fixture +def installed(monkeypatch): + saved = (gen_mod.ResponseGenerator._make_thinking_budget_criteria, + ar.BatchGenerator.insert, ar.GenerationBatch._step, + ar.PromptProcessingBatch.generate, + ar.SpeculativeGenerationBatch.next) + monkeypatch.setattr( + gen_mod.ResponseGenerator, "_make_thinking_budget_criteria", + lambda self, args, input_ids: None) + monkeypatch.setattr( + ar.BatchGenerator, "insert", + lambda self, prompts, **kw: list(range(100, 100 + len(prompts)))) + monkeypatch.setattr(ar.GenerationBatch, "_step", + lambda self: self.sampler._kq_rows) + monkeypatch.setattr(ar.PromptProcessingBatch, "generate", + lambda self, sampler, *a, **k: sampler._kq_rows) + monkeypatch.setattr(ar.SpeculativeGenerationBatch, "next", + lambda self: self.sampler._kq_rows) + sr._PENDING.clear() + sr.install_per_request_seed() + yield + (gen_mod.ResponseGenerator._make_thinking_budget_criteria, + ar.BatchGenerator.insert, ar.GenerationBatch._step, + ar.PromptProcessingBatch.generate, + ar.SpeculativeGenerationBatch.next) = saved + + +def test_insert_registers_the_request_seed(installed): + s = _sampler() + rg = SimpleNamespace() + gen_mod.ResponseGenerator._make_thinking_budget_criteria( + rg, SimpleNamespace(seed=42, temperature=1.0), None) + bg = SimpleNamespace(sampler=s) + uids = ar.BatchGenerator.insert(bg, [[1, 2, 3]]) + assert s._kq_row_seeds == {uids[0]: 42} + + +def test_greedy_request_seed_is_ignored(installed): + s = _sampler() + rg = SimpleNamespace() + gen_mod.ResponseGenerator._make_thinking_budget_criteria( + rg, SimpleNamespace(seed=42, temperature=0), None) + bg = SimpleNamespace(sampler=s) + ar.BatchGenerator.insert(bg, [[1, 2, 3]]) + assert s._kq_row_seeds == {} + + +def test_unseeded_insert_registers_nothing(installed): + s = _sampler() + rg = SimpleNamespace() + gen_mod.ResponseGenerator._make_thinking_budget_criteria( + rg, SimpleNamespace(seed=None, temperature=1.0), None) + bg = SimpleNamespace(sampler=s) + ar.BatchGenerator.insert(bg, [[1, 2, 3]]) + assert s._kq_row_seeds == {} + + +def test_step_publishes_uids_only_when_seeds_exist(installed): + s = _sampler() + gb = SimpleNamespace(sampler=s, uids=[1, 2]) + assert ar.GenerationBatch._step(gb) is None # no seeds: no context + s._kq_row_seeds[1] = 9 + assert ar.GenerationBatch._step(gb) == [1, 2] + assert s._kq_rows is None # cleared after the step + + +def test_prompt_generate_and_spec_next_publish_uids(installed): + s = _sampler() + s._kq_row_seeds[1] = 9 + pb = SimpleNamespace(uids=[1]) + assert ar.PromptProcessingBatch.generate(pb, s) == [1] + sb = SimpleNamespace(sampler=s, _all_uids=[1, 3]) + assert ar.SpeculativeGenerationBatch.next(sb) == [1, 3] + assert s._kq_rows is None diff --git a/tests/test_spec_kv_quant.py b/tests/test_spec_kv_quant.py index e13dbe7..c72371d 100644 --- a/tests/test_spec_kv_quant.py +++ b/tests/test_spec_kv_quant.py @@ -132,6 +132,34 @@ def test_quantized_trim_rollback_exact(): assert mx.array_equal(pa, pb).item() +def test_prefix_cache_quantized_roundtrip(): + # The spec prefix cache must snapshot quantized targets: state entries + # are (packed, scales, biases) triples, and the old array-only path + # crashed every MTP store under KV_BITS. Store, mutate the live cache, + # restore into a fresh one, and the pre-mutation triples must match. + from gmlx.prefix_cache import SpecPrefixCache + + mx.random.seed(11) + a = QuantizedKVCache(group_size=64, bits=8) + _fill(a, [mx.random.normal((1, 2, 40, 64)).astype(mx.bfloat16)]) + ref = [tuple(mx.contiguous(x) for x in side) for side in a.state] + cache = SpecPrefixCache() + cache.store(mx.array(list(range(40))), [a], mx.zeros((1, 1, 8))) + _fill(a, [mx.random.normal((1, 2, 5, 64)).astype(mx.bfloat16)]) + + hit = cache.lookup(mx.array(list(range(41)))) + assert hit is not None and hit[0] == 40 + b = QuantizedKVCache(group_size=64, bits=8) + cache.restore(hit[1], [b]) + assert b.offset == 40 + for got, want in zip((b.keys, b.values), ref): + for pg, pw in zip(got, want): + assert mx.array_equal(pg, pw).item() + # restored cache must keep decoding past the snapshot point + b.update_and_fetch(*2 * (mx.random.normal((1, 2, 1, 64)).astype(mx.bfloat16),)) + assert b.offset == 41 + + def _dequant_ref(q, qc, scale): # upstream per-token verify loop on dequantized KV keys = mx.dequantize(*qc.keys, group_size=qc.group_size, bits=qc.bits)