diff --git a/CHANGELOG.md b/CHANGELOG.md index 2db3075..f85ad94 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,12 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### Added +- decode_prefill_ratio accepts "auto" and it is the new default: live + streams keep at least half their decode rate while deep prompts admit, + and pacing stands down wherever it would not help (simultaneous + bursts, cheap chunks, stuck queues). A numeric ratio pins the previous + static behavior; GMLX_DECODE_PREFILL_AUTO=0 reverts on a live server. + - GMLX_SERVE_MEMSTATS=path.jsonl writes a per-tick serve memory trace: MLX counters, free-headroom estimate, and per-owner cache byte attribution with allocation shapes marked on change, for diagnosing diff --git a/docs/performance.md b/docs/performance.md index d797b98..b8170e5 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -278,22 +278,45 @@ What needs managing is admission: a new request's prompt must prefill while existing streams are mid-decode. Prefill runs in 2048-token chunks, and a scheduler that simply alternates one decode step with one chunk lets a long admission starve live streams, because at depth a chunk costs hundreds of -decode steps' worth of GPU time. The server paces admissions instead: -`decode_prefill_ratio` (default `1.0`) admits the next chunk only after the -decode batch has received that multiple of the previous chunk's GPU time. At -the default, live streams keep roughly half their throughput while a prompt -is admitted, and the incoming request's time-to-first-token stretches by up -to (1 + ratio)x under load. Raise the ratio when live-stream decode matters -most, lower it toward `0` when time-to-first-token does; `0` restores strict -alternation. Prefill runs at full speed whenever nothing is decoding, so -single-client serving is unaffected. +decode steps' worth of GPU time. Whether pacing admissions helps is decided +by that same quantity: when a chunk costs a live stream many decode steps +(deep context), stock scheduling starves it and pacing rescues it; when +chunks are cheap (shallow prompts, warm prefix hits), pacing only delays +admission, and a delayed admission narrows the decode batch that aggregate +throughput comes from. + +`decode_prefill_ratio` (default `auto`) measures this per tick and paces +only when an already-decoding stream that was admitted before the waiters +arrived would otherwise fall below half its batched decode rate. For +simultaneous bursts (no incumbent to protect), cheap chunks, and queued +waiters held behind paced admissions past a deadline it runs stock +scheduling, so one setting serves shallow-burst and deep-second-client +load alike. Paced admission bounds every waiter's time-to-first-token at +twice its unpaced prefill, even when several arrive at once. The +deadline counts only time pacing itself is responsible for: a waiter +blocked by a full decode batch or by the memory admission gate is not +aging toward it, since running unpaced would not admit that waiter any +sooner. + +A numeric value pins the static behavior: the decode batch receives that +multiple of each chunk's GPU time before the next chunk is admitted, and at +`1.0` live streams keep roughly half their throughput while a prompt is +admitted. `0` restores strict alternation. Static pacing has two costs +worth naming. A waiter's time-to-first-token stretch compounds with queue +depth, since each waiter also waits out the throttled prefill of everyone +ahead of it: several-fold at moderate bursts, not the single-admission +(1 + ratio)x. And delaying admission keeps the decode batch narrow, which +at burst concurrency can cost aggregate throughput outright. Prefill runs +at full speed whenever nothing is decoding, so single-client serving is +unaffected under every setting. The deeper the context, the more this matters. In our serve benchmarks on -the same 35B-A3B, adding a second client at 14k tokens used to drop -aggregate decode to 0.57x of single-stream; paced, it lands above -single-stream. At 50k tokens each of two streams held ~10 tok/s under -alternation and ~50 tok/s paced, because a 50k admission previously froze -live streams for tens of seconds. The key is `server.decode_prefill_ratio` +the same 35B-A3B, a second client arriving at 14k tokens under strict +alternation froze the live stream to 4 percent of its decode rate for the +whole admission; paced, it keeps 80 percent, with the second client's +time-to-first-token unchanged. At 50k tokens the admission is roughly a +minute of prefill and the live stream holds 54 percent instead of 3, a +~26x higher rate through the window. The key is `server.decode_prefill_ratio` ([server-config.md](server-config.md)), the `serve` flag is `--decode-prefill-ratio`, and the `GMLX_DECODE_PREFILL_RATIO` env is read per scheduler tick, so it can be changed on a live server. diff --git a/docs/server-config.md b/docs/server-config.md index 53a27d9..bf2c91d 100644 --- a/docs/server-config.md +++ b/docs/server-config.md @@ -181,7 +181,7 @@ server: # (null => mlx-vlm's own default, 600; 0 => never) prefill_step_size: null # prefill chunk size in tokens for every model on this server # (null => the default, 2048; lower caps peak memory on long prompts) - decode_prefill_ratio: null # decode GPU-time share per admission prefill chunk under load + decode_prefill_ratio: null # admission pacing: auto (default) or a static GPU-time share # (null => the default, 1.0; 0 => strict alternation; see below) prefill_tick_ms: null # wall-clock budget per prefill chunk while streams decode; # chunks are halved to fit (null => the default, 500; 0 => full chunks) @@ -247,14 +247,22 @@ closed, so it cannot be a per-model `load:` key. Also available as `decode_prefill_ratio` paces admission prefills against live decode. Stock scheduling runs one decode step per prefill chunk, so while any request prefills, every decoding stream advances ~1 token per chunk -- at deep context -that is a multi-second stall per admission. With pacing (default `1.0`), a -prefill chunk is admitted only after the decode batch has received that -multiple of the chunk's GPU time: live streams keep ~half throughput during -admissions, and the incoming request's time-to-first-token stretches up to -~(1+ratio)x while decode is busy. Raise the ratio to favor decode further, -lower it toward `0` for TTFT-critical serving, `0` restores stock scheduling. -Prefill runs at full speed whenever nothing is decoding, so a single-stream -server is unaffected. Also available as `--decode-prefill-ratio` on `serve` +that is a multi-second stall per admission. The default `auto` paces only +when an already-decoding stream admitted before the waiters would otherwise +fall below half its batched decode rate (the floor; +`GMLX_DECODE_PREFILL_FLOOR`), and runs stock scheduling for simultaneous +bursts, cheap chunks, and queued waiters held behind paced admissions +past a deadline (a prompt already being prefilled is bounded by pacing +itself, and time blocked by capacity rather than pacing does not age +toward the deadline). A numeric +value pins static pacing: a prefill chunk is admitted only after the decode +batch has received that multiple of the chunk's GPU time; live streams then +keep ~half throughput during admissions at `1.0`, while a waiter's +time-to-first-token stretch compounds with queue depth (each waiter also +waits out the throttled prefill of everyone ahead of it) and delayed +admission narrows the decode batch. `0` restores stock scheduling. Prefill +runs at full speed whenever nothing is decoding, so a single-stream server +is unaffected. Also available as `--decode-prefill-ratio` on `serve` (the flag wins over the config) or an exported `GMLX_DECODE_PREFILL_RATIO` (read per scheduler tick, so it can be flipped on a live server). Applies to speculative (MTP) serving too. Background and measured effects: diff --git a/gmlx/auto_ratio.py b/gmlx/auto_ratio.py new file mode 100644 index 0000000..93e30ac --- /dev/null +++ b/gmlx/auto_ratio.py @@ -0,0 +1,349 @@ +"""Dynamic decode-prefill pacing (``decode_prefill_ratio: auto``). + +The static pacing ratio cannot be right at both ends of the depth range. +Pacing exists to keep a live decoding stream from being starved by a deep +prompt's prefill, and how much stock scheduling costs that stream is +decided by one measurable quantity: C, the wall cost of one prefill chunk +expressed in decode steps. At depth C is large and an unpaced prefill +puts the live stream near 1/(1+C) of its rate, so pacing rescues it. On +cheap chunks (shallow prompts, warm prefix hits, small models) stock +already leaves the stream above any reasonable floor and pacing only +delays admission, narrowing the decode batch that aggregate throughput +comes from. + +Auto mode enforces one user-facing constraint, the retention floor rho +(default 0.5): a live decoding row is not pushed below rho of its +contemporaneous no-prefill batched rate. The floor is about sustained +starvation, so pacing shapes multi-chunk prefill trains; C is measured, +not predicted, so a single-chunk admission (a warm-prefix suffix, a +short prompt) completes before its cost is observable, and its worst +case for an incumbent is one chunk of stall, bounded by the prefill +step size. Steady-interleave algebra gives +retention r/(1+r) at ratio r independent of depth, so the floor fixes the +paced ratio at rho/(1-rho) and the chunk-cost threshold at (1-rho)/rho; +there is no useful middle ratio (a mid ratio pays for both a narrow +decode batch and a stretched prefill). Auto therefore only ever selects +the paced ratio or 0, per tick: + + r = rho/(1-rho) if an incumbent row exists (admitted at least + grace_ms before the oldest waiter arrived) + and C > (1-rho)/rho (hysteresis + dwell) + and no queued waiter's paced wait exceeds deadline_s + 0 otherwise + +A waiter is a queued sequence or a row of the active prompt batch: a +prompt being chunked is still competing prefill work, and dropping +pacing at promotion would unpace the whole chunk train one tick after +it starts. Only queued waiters age toward the deadline; a prefilling +waiter's TTFT is already bounded multiplicatively by the paced ratio +(its chunk train stretches by at most (1+r)x), while a queued waiter's +wait is what pacing cannot bound. + +The deadline ages pacing-attributable seconds, not wall time since +arrival: a queued waiter accrues age only on ticks where a prompt +train is live and the previous tick resolved paced, bounded at 1 s per +tick. Time blocked by capacity is deliberately excluded. A waiter +behind a full decode batch or a memory-gate decline is not waiting on +pacing, and abandoning the floor for it cannot admit it any sooner: +ratio 0 creates neither batch slots nor headroom, so an arrival-age +deadline would starve every incumbent to buy nothing. The sustained +regime C run measured exactly this shape: four waiters capacity-held +~50 s at saturation whose final admission was correctly paced, while +waiters genuinely queued behind paced trains aged to the deadline and +were shed unpaced. + +Only the C conjunct enforces the floor; incumbency decides whether anyone +is owed it. The deadline conjunct knowingly abandons the floor in favor +of waiters, and every tick that does so logs the abandonment with the +incumbent's projected retention. + +The deadline is the only abandonment policy; there is deliberately no +queue-depth term. Waiters drain into the prompt batch within a tick or +two of arriving, so a count of queued sequences is a race against +promotion, not a measure of pressure (the same burst can read 4 or 0 +depending on tick phase). Real pressure is waiters that cannot promote, +and those age in the queue until the deadline abandons the floor for +them. Standing down on count alone was measured to freeze the incumbent +to ~0.02 retention during a burst's batched prefill while also finishing +the incumbent later than pacing would have; the paced alternative holds +the floor and bounds every burst waiter's TTFT at (1+r)x unpaced. + +Signals are measured in the scheduler wrapper with no new sync: the +decode step cost s is an exponentially weighted mean per decode width +(width changes exactly at admission), the first decode tick after a chunk +is skipped (its bracket reads the chunk's sync, not a step), and C is +computed residue-corrected as max(0, last_chunk - s) / s because the +chunk bracket absorbs up to one in-flight decode step. Waiter arrival and +row admission stamps are kept per uid, bounded by one tick. + +Admission-gate coupling: on a declined tick the gate hides the pending +list from the stock body, so this resolver sees no queued waiters and +accrues nothing, which is the exclusion the design wants (the gate, not +pacing, is why they wait). The one thing the resolver must do is keep +stamps and accruals alive for hidden waiters instead of pruning them as +departed, so it unions the gate's deferred-uid dict +(``_kq_admit_deferred_s``, read through getattr defaults so either +feature works without the other) into the waiter set for retention. + +State lives on the generator under ``_kq_auto_`` attributes. The kill +switch GMLX_DECODE_PREFILL_AUTO=0 resolves auto to the static paced ratio +without a restart; any numeric ratio at any precedence layer bypasses +auto entirely. + +Calibration knobs (documented here, not in the public config reference): + + GMLX_DECODE_PREFILL_FLOOR retention floor rho (default 0.5) + GMLX_DECODE_PREFILL_HYST C hysteresis band (default 1.5) + GMLX_DECODE_PREFILL_DWELL_MS min time in a C state (default 1000) + GMLX_DECODE_PREFILL_DEADLINE_S queued paced-wait forcing ratio 0 + (default 10) + GMLX_DECODE_PREFILL_GRACE_MS incumbency grace (default 500) + GMLX_DECODE_PREFILL_ALPHA step-cost smoothing (default 0.2) + GMLX_DECODE_PREFILL_LOG_S transition log rate limit (default 1; + 0 logs every transition, for + debugging sub-second episodes) +""" + +from __future__ import annotations + +import logging +import os +import time + +_log = logging.getLogger(__name__) + +_LOG_EVERY_S = 1.0 + + +def _envf(name: str, default: float) -> float: + try: + return float(os.environ.get(name, "") or default) + except ValueError: + return default + + +def floor_rho() -> float: + rho = _envf("GMLX_DECODE_PREFILL_FLOOR", 0.5) + return min(max(rho, 0.05), 0.95) + + +def paced_ratio() -> float: + rho = floor_rho() + return rho / (1.0 - rho) + + +def c_threshold() -> float: + rho = floor_rho() + return (1.0 - rho) / rho + + +def deadline_s() -> float: + return _envf("GMLX_DECODE_PREFILL_DEADLINE_S", 10.0) + + +class _AutoState: + """Per-generator resolver state (held under gen._kq_auto).""" + + def __init__(self): + self.s_by_width: dict[int, float] = {} + self.s_samples: dict[int, int] = {} + self.skip_next_step = False + self.first_seen: dict = {} + self.admitted_at: dict = {} + self.paced_wait: dict = {} + self.c_on: bool | None = None + self.c_since = 0.0 + self.last_resolve_t: float | None = None + self.last_log = 0.0 + self.last_logged: float | None = None + self.last_resolved: float | None = None + + +def _state(gen) -> _AutoState: + st = getattr(gen, "_kq_auto", None) + if st is None: + st = gen._kq_auto = _AutoState() + return st + + +def observe(gen, dt: float, prompt_delta: float, rows: int) -> None: + """Fold one tick's bracket into the step-cost estimate. A chunk tick + poisons the next decode bracket (dispatched pre-chunk, completed in + the chunk's sync), so it is skipped.""" + st = _state(gen) + if prompt_delta > 0.0: + st.skip_next_step = True + return + if rows <= 0 or dt <= 0.0: + return + if st.skip_next_step: + st.skip_next_step = False + return + alpha = min(max(_envf("GMLX_DECODE_PREFILL_ALPHA", 0.2), 0.01), 1.0) + prev = st.s_by_width.get(rows) + real = st.s_samples.get(rows, 0) + st.s_by_width[rows] = dt if (prev is None or real == 0) else ( + (1 - alpha) * prev + alpha * dt) + st.s_samples[rows] = real + 1 + + +def _step_cost(st: _AutoState, width: int) -> tuple[float | None, bool]: + """(s for this width, frozen). An unfed bucket seeds from the nearest + populated width; the seed is exactly the stale value at the moment of + admission, so the C state holds (frozen) until real samples land.""" + if st.s_samples.get(width, 0) > 0: + return st.s_by_width[width], False + fed = [w for w, n in st.s_samples.items() if n > 0] + if not fed: + return None, True + nearest = min(fed, key=lambda w: abs(w - width)) + st.s_by_width[width] = st.s_by_width[nearest] + return st.s_by_width[nearest], True + + +def _pb_uids(gen) -> list: + """Rows of the active prompt batch: promoted waiters mid-prefill.""" + pb = gen._prompt_batch + return list(getattr(pb, "uids", ())) if pb is not None else [] + + +def _stamp(gen, st: _AutoState, now: float) -> None: + # A waiter keeps its arrival stamp across promotion into the prompt + # batch; the stamp dies when its rows reach the decode batch (or the + # request is gone). Dropping it at promotion would both end pacing a + # tick after the chunk train starts and reset incumbency comparisons. + # Gate-deferred waiters are hidden from the pending list on declined + # ticks; their stamps and accruals are retained, not pruned. + pending_uids = [s[0] for s in gen._unprocessed_sequences] + gate_hidden = list(getattr(gen, "_kq_admit_deferred_s", None) or {}) + waiter_uids = pending_uids + _pb_uids(gen) + gate_hidden + for uid in waiter_uids: + st.first_seen.setdefault(uid, now) + keep = set(waiter_uids) + for uid in list(st.first_seen): + if uid not in keep: + del st.first_seen[uid] + for uid in list(st.paced_wait): + if uid not in keep: + del st.paced_wait[uid] + live = list(getattr(gen._generation_batch, "uids", ())) + for uid in live: + st.admitted_at.setdefault(uid, now) + for uid in list(st.admitted_at): + if uid not in set(live): + del st.admitted_at[uid] + + +def _c_term(gen, st: _AutoState, now: float) -> tuple[bool, float]: + """C-threshold conjunct with hysteresis, dwell, and the width-change + freeze. Returns (pacing_wanted, C).""" + width = len(gen._generation_batch) + s, frozen = _step_cost(st, width) + last_chunk = getattr(gen, "_kq_last_chunk_time", 0.0) + if s is None or s <= 0.0 or last_chunk <= 0.0: + return (bool(st.c_on), 0.0) + c = max(0.0, last_chunk - s) / s + if frozen and st.c_on is not None: + return st.c_on, c + c_on = c_threshold() + hyst = max(_envf("GMLX_DECODE_PREFILL_HYST", 1.5), 1.0) + dwell = _envf("GMLX_DECODE_PREFILL_DWELL_MS", 1000.0) / 1e3 + if st.c_on is None: + st.c_on = c > c_on + st.c_since = now + return st.c_on, c + dwelled = (now - st.c_since) >= dwell + if not st.c_on and c > c_on and dwelled: + st.c_on, st.c_since = True, now + elif st.c_on and c < c_on / hyst and dwelled: + st.c_on, st.c_since = False, now + return st.c_on, c + + +def resolve(gen, now: float | None = None) -> float: + """The effective ratio for this tick under auto mode.""" + if os.environ.get("GMLX_DECODE_PREFILL_AUTO", "1") == "0": + return paced_ratio() + st = _state(gen) + if now is None: + now = time.perf_counter() + _stamp(gen, st, now) + pending = gen._unprocessed_sequences + prefilling = _pb_uids(gen) + + # Queued waiters age only while behind a live paced train; capacity + # waits (no train, or gate-hidden pending) accrue nothing, since + # ratio 0 cannot admit them sooner. The per-tick bound keeps a + # stalled tick from charging its gap to pacing. + dt = min(max(now - st.last_resolve_t, 0.0), 1.0) \ + if st.last_resolve_t is not None else 0.0 + st.last_resolve_t = now + if dt > 0.0 and prefilling and (st.last_resolved or 0.0) > 0.0: + for s in pending: + st.paced_wait[s[0]] = st.paced_wait.get(s[0], 0.0) + dt + + if not pending and not prefilling: + return _resolved(gen, st, 0.0, now, "no waiters") + + # Incumbency compares against the oldest waiter of either kind; + # the deadline (below) ages only the queued ones. + oldest_seen = None + queued_uid, queued_wait = None, 0.0 + for uid in [s[0] for s in pending] + prefilling: + seen = st.first_seen.get(uid, now) + if oldest_seen is None or seen < oldest_seen: + oldest_seen = seen + for s in pending: + wait = st.paced_wait.get(s[0], 0.0) + if queued_uid is None or wait > queued_wait: + queued_uid, queued_wait = s[0], wait + + grace = _envf("GMLX_DECODE_PREFILL_GRACE_MS", 500.0) / 1e3 + incumbent = None + for uid, adm in st.admitted_at.items(): + if adm + grace < (oldest_seen if oldest_seen is not None else now): + incumbent = uid + break + if incumbent is None: + return _resolved( + gen, st, 0.0, now, + f"no incumbent (waiting={len(pending)}, " + f"width={len(gen._generation_batch)})") + + c_wants, c = _c_term(gen, st, now) + if not c_wants: + return _resolved(gen, st, 0.0, now, + f"chunk {c:.1f} steps <= {c_threshold():.1f}") + + # Deadline: queued waiters only. A prefilling waiter's TTFT is + # bounded by the paced ratio itself (chunk train stretches at most + # (1+r)x); a queued waiter's paced wait has no such bound, so it is + # the one the floor is abandoned for. + deadline = deadline_s() + if queued_uid is not None and queued_wait > deadline: + return _resolved( + gen, st, 0.0, now, + f"queued waiter paced {queued_wait:.1f}s > deadline " + f"{deadline:.1f}s (floor abandoned: incumbent uid={incumbent} " + f"projected retention {1.0 / (1.0 + c):.2f})") + + return _resolved( + gen, st, paced_ratio(), now, + f"incumbent uid={incumbent} (admitted " + f"{now - st.admitted_at[incumbent]:.1f}s ago), chunk {c:.1f} " + f"steps, waiting={len(pending)}+{len(prefilling)}") + + +def _resolved(gen, st: _AutoState, ratio: float, now: float, + reason: str) -> float: + # Compare against last_logged, not last_resolved: a rate-limited + # transition must still log once the window opens, or a persisted + # state stays unlogged. + every = _envf("GMLX_DECODE_PREFILL_LOG_S", _LOG_EVERY_S) + if ratio != st.last_logged and now - st.last_log >= every: + st.last_log = now + st.last_logged = ratio + _log.info("[sched] pacing %s: %s", + "on" if ratio > 0 else "off", reason) + st.last_resolved = ratio + return ratio diff --git a/gmlx/batch_sched.py b/gmlx/batch_sched.py index 9dc7691..7a4e109 100644 --- a/gmlx/batch_sched.py +++ b/gmlx/batch_sched.py @@ -32,9 +32,14 @@ Ratio resolution: `gmlx serve --decode-prefill-ratio` / config `server.decode_prefill_ratio` both export GMLX_DECODE_PREFILL_RATIO; the wrapper reads the env per tick (install order never matters, and a live -server can be re-paced for an A/B). 0 disables (stock 1:1). Very large -ratios starve admission at depth: each pending chunk waits ratio x -chunk_time of decode -- TTFT of queued requests stretches accordingly. +server can be re-paced for an A/B). The default is ``auto``: the +per-tick ratio comes from ``auto_ratio``, which selects the paced ratio +or 0 from measured signals. A numeric value pins a static split +(0 disables: stock 1:1); every other mechanism below (owed-decode +arithmetic, stash-and-restore, pressure notes) is identical under both +spellings. Very large static ratios starve admission at depth: each +pending chunk waits ratio x chunk_time of decode -- TTFT of queued +requests stretches accordingly. """ from __future__ import annotations @@ -48,22 +53,28 @@ _INSTALLED_FLAG = "_kq_gguf_decode_priority_sched" -_ratio_memo = ("", 1.0) +_ratio_memo: tuple[str, tuple[str, float | None]] = ("", ("static", 1.0)) -def _ratio() -> float: - """Parse GMLX_DECODE_PREFILL_RATIO, memoized on the raw string; warn - once per bad value rather than silently defaulting.""" +def _ratio() -> tuple[str, float | None]: + """Parse GMLX_DECODE_PREFILL_RATIO into ``(mode, value)``: + ``("static", r)`` or ``("auto", None)``. Memoized on the raw string + (in auto mode only the parse is memoized, never a resolved ratio); + warns once per bad value rather than silently defaulting.""" global _ratio_memo - raw = os.environ.get("GMLX_DECODE_PREFILL_RATIO", "1.0") + raw = os.environ.get("GMLX_DECODE_PREFILL_RATIO", "auto") if raw == _ratio_memo[0]: return _ratio_memo[1] - try: - val = float(raw) - except ValueError: - _log.warning( - "GMLX_DECODE_PREFILL_RATIO=%r is not a number; using 1.0", raw) - val = 1.0 + if raw.strip().lower() == "auto": + val: tuple[str, float | None] = ("auto", None) + else: + try: + val = ("static", float(raw)) + except ValueError: + _log.warning( + "GMLX_DECODE_PREFILL_RATIO=%r is not a number or 'auto'; " + "using auto", raw) + val = ("auto", None) _ratio_memo = (raw, val) return val @@ -96,10 +107,10 @@ def _observed_next(self, **kwargs): _pd.note_chunk_cost(spent) return out - def _paced_next(self, **kwargs): - # Pace only when decode AND prefill work are both live; otherwise - # stock behavior (prefill at full speed, untouched TTFT). - ratio = _ratio() + def _tick(self, ratio, **kwargs): + # One scheduler tick at a resolved ratio. Pace only when decode + # AND prefill work are both live; otherwise stock behavior + # (prefill at full speed, untouched TTFT). if ratio <= 0: # Ratio 0 is documented as stock: clear the tick term's # decode-pressure reading (0 clears, per the hook contract) @@ -147,8 +158,29 @@ def _paced_next(self, **kwargs): # observe the chunk cost. return _observed_next(self, **kwargs) + def _paced_next(self, **kwargs): + mode, value = _ratio() + if mode != "auto": + return _tick(self, value, **kwargs) + from . import auto_ratio + + ratio = auto_ratio.resolve(self) + rows = len(self._generation_batch) + before = self._prompt_time_counter + tic = time.perf_counter() + try: + return _tick(self, ratio, **kwargs) + finally: + # The bracket rides syncs the tick already performs (chunk + # mx.eval / decode blocking on the previous step); auto's + # observe skips the poisoned first bracket after a chunk. + auto_ratio.observe( + self, time.perf_counter() - tic, + self._prompt_time_counter - before, rows) + setattr(_paced_next, _INSTALLED_FLAG, True) _ar.BatchGenerator._next = _paced_next + mode, value = _ratio() _log.info( - "decode-priority prefill pacing installed (ratio=%.2f)", _ratio() - ) + "decode-priority prefill pacing installed (%s)", + "auto" if mode == "auto" else f"ratio={value:.2f}") diff --git a/gmlx/config.py b/gmlx/config.py index cd03b11..871f248 100644 --- a/gmlx/config.py +++ b/gmlx/config.py @@ -417,8 +417,9 @@ class ServerCfg: # Decode-priority prefill pacing ratio: a live decode batch gets this # multiple of each prefill chunk's GPU time before the next chunk is # admitted (1.0 ~= 50/50 split; 0 = stock 1 decode step : 1 chunk). - # None => leave the env / branch default (1.0) in place. - decode_prefill_ratio: float | None = None + # None => leave the env / branch default in place, which is "auto" + # (dynamic pacing via auto_ratio); a number pins a static split. + decode_prefill_ratio: float | str | None = None # Prefill tick budget in wall-clock ms: while decode rows are live, each # prefill chunk is halved until its predicted wall time (from the last # observed chunk cost) fits this budget, bounding the per-chunk decode @@ -1246,6 +1247,14 @@ def _normalize_cache(where: str, raw) -> dict: return cache +def _coerce_ratio(key: str, v, *, where: str = "server"): + """decode_prefill_ratio: a float, or the literal string "auto" + (case-insensitive, surrounding whitespace stripped).""" + if isinstance(v, str) and v.strip().lower() == "auto": + return "auto" + return _coerce_num(key, v, float, where=where) + + def _coerce_num(key: str, v, cast, *, where: str = "server"): """Coerce a numeric config key (YAML may carry it quoted as a string), raising a ConfigError naming the key and the bad value. ``None`` passes.""" @@ -1656,8 +1665,8 @@ def build_config(doc: dict) -> ServerCfg: "token_queue_timeout_s", srv.get("token_queue_timeout_s"), float), prefill_step_size=_coerce_num( "prefill_step_size", srv.get("prefill_step_size"), int), - decode_prefill_ratio=_coerce_num( - "decode_prefill_ratio", srv.get("decode_prefill_ratio"), float), + decode_prefill_ratio=_coerce_ratio( + "decode_prefill_ratio", srv.get("decode_prefill_ratio")), prefill_tick_ms=_coerce_num( "prefill_tick_ms", srv.get("prefill_tick_ms"), float), cache_limit_gb=_coerce_num( diff --git a/gmlx/server.py b/gmlx/server.py index 384d1b9..e520fca 100644 --- a/gmlx/server.py +++ b/gmlx/server.py @@ -59,6 +59,17 @@ _DEFAULT_DISCOVER_DIR = "." # zero-config bare start scans the cwd +def _ratio_flag(raw: str): + """--decode-prefill-ratio value: a float or the literal 'auto'.""" + if raw.strip().lower() == "auto": + return "auto" + try: + return float(raw) + except ValueError: + raise argparse.ArgumentTypeError( + f"expected a number or 'auto', got {raw!r}") + + def _has_uvloop() -> bool: try: import uvloop # noqa: F401 @@ -780,11 +791,13 @@ def _add_serve_args(ap: argparse.ArgumentParser) -> None: "it to cap peak memory on long prompts, at some " "prefill-throughput cost. Also via PREFILL_STEP_SIZE; " "config mode: server.prefill_step_size.") - ap.add_argument("--decode-prefill-ratio", type=float, default=None, - metavar="R", + ap.add_argument("--decode-prefill-ratio", type=_ratio_flag, + default=None, metavar="R", help="Decode GPU-time share per prefill chunk under load " - "(default 1.0 ~= 50/50; 0 = stock scheduling). Also " - "via GMLX_DECODE_PREFILL_RATIO; config mode: " + "(default auto: paced or stock per tick from " + "measured load; numeric pins a static split, " + "1.0 ~= 50/50, 0 = stock scheduling). Also via " + "GMLX_DECODE_PREFILL_RATIO; config mode: " "server.decode_prefill_ratio.") ap.add_argument("--prefill-tick-ms", type=float, default=None, metavar="MS", @@ -1514,7 +1527,15 @@ def _serve(cfg: ServerCfg, a, reload_fn) -> int: if ratio is None: ratio = cfg.decode_prefill_ratio if ratio is not None: - if ratio < 0: + if isinstance(ratio, str): + from .auto_ratio import (c_threshold, deadline_s, floor_rho, + paced_ratio) + os.environ["GMLX_DECODE_PREFILL_RATIO"] = "auto" + print(f"[server] decode-prefill pacing ratio: auto " + f"(floor {floor_rho():.2f} -> paced {paced_ratio():.2f}, " + f"chunk threshold {c_threshold():.1f}, " + f"deadline {deadline_s():.1f}s)") + elif ratio < 0: print(f"[server] ignoring negative decode-prefill ratio {ratio}") else: os.environ["GMLX_DECODE_PREFILL_RATIO"] = str(ratio) diff --git a/tests/test_auto_ratio.py b/tests/test_auto_ratio.py new file mode 100644 index 0000000..90505d6 --- /dev/null +++ b/tests/test_auto_ratio.py @@ -0,0 +1,276 @@ +"""Auto pacing resolver: one case per conjunct, hysteresis, dwell, grace, +deadline accrual, gate coupling. Pure logic over a fake generator with an +injected clock (resolve takes ``now``; nothing sleeps).""" + +import gmlx.auto_ratio as ar + + +class FakeBatch: + def __init__(self, uids): + self.uids = list(uids) + + def __len__(self): + return len(self.uids) + + +class FakeGen: + def __init__(self, rows=(101,), pending=()): + self._generation_batch = FakeBatch(rows) + self._prompt_batch = None + self._unprocessed_sequences = [ + (u, [0] * 10, 64, {}, None, None) for u in pending] + self._kq_last_chunk_time = 0.0 + self._prompt_time_counter = 0.0 + + +def _feed_step(g, s=0.05, width=1): + ar.observe(g, s, 0.0, width) + + +def _incumbent_gen(now=0.0): + """Row 101 admitted at t=now, step cost fed, expensive chunk seen.""" + g = FakeGen(rows=(101,), pending=()) + ar.resolve(g, now) # stamps admitted_at[101]=now + _feed_step(g) + g._kq_last_chunk_time = 1.0 # C = (1.0 - 0.05) / 0.05 = 19 + return g + + +def _add_waiter(g, uid=7): + g._unprocessed_sequences.append((uid, [0] * 10, 64, {}, None, None)) + + +def test_no_waiters_resolves_zero(): + g = _incumbent_gen() + assert ar.resolve(g, 1.0) == 0.0 + + +def test_incumbent_plus_costly_chunk_paces(): + g = _incumbent_gen(0.0) + _add_waiter(g) + assert ar.resolve(g, 1.0) == ar.paced_ratio() + + +def test_burst_has_no_incumbent(): + # every row admitted the same tick the waiters were first seen + g = FakeGen(rows=(1, 2), pending=(3, 4)) + _feed_step(g, width=2) + g._kq_last_chunk_time = 1.0 + assert ar.resolve(g, 5.0) == 0.0 # tie rule: same-tick stamp loses + + +def test_grace_blocks_fresh_incumbent(): + g = FakeGen(rows=(101,)) + ar.resolve(g, 0.0) + _feed_step(g) + g._kq_last_chunk_time = 1.0 + _add_waiter(g) + # waiter arrives 0.3s after admission, inside the 500ms grace + assert ar.resolve(g, 0.3) == 0.0 + # the comparison is fixed at arrival: the same waiter never grants + # incumbency however long both wait + assert ar.resolve(g, 3.0) == 0.0 + # a waiter that arrives later (after the first was admitted) does + # see the long-lived row as an incumbent + g._unprocessed_sequences.clear() + ar.resolve(g, 3.5) + _add_waiter(g, 8) + assert ar.resolve(g, 4.0) == ar.paced_ratio() + + +def test_cheap_chunk_stays_stock(): + g = _incumbent_gen(0.0) + g._kq_last_chunk_time = 0.08 # C = (0.08 - 0.05) / 0.05 = 0.6 <= 1 + _add_waiter(g) + assert ar.resolve(g, 1.0) == 0.0 + + +def test_c_hysteresis_and_dwell(): + g = _incumbent_gen(0.0) + _add_waiter(g) + assert ar.resolve(g, 1.0) == ar.paced_ratio() # C=19, state on + # C falls inside the band (< C_on but > C_on/hyst): stays on + g._kq_last_chunk_time = 0.09 # C = 0.8 + assert ar.resolve(g, 2.5) == ar.paced_ratio() + # C below C_on/hyst but dwell not elapsed since the flip at t=1.0 + g._kq_last_chunk_time = 0.06 # C = 0.2 < 1/1.5 + assert ar.resolve(g, 1.5) == ar.paced_ratio() + # dwell elapsed: flips off + assert ar.resolve(g, 2.6) == 0.0 + + +def test_burst_stays_paced(): + # The regime C burst regression: a burst of queued waiters is not + # queue pressure (they promote into the prompt batch within a tick + # or two), and standing down on count was a race against promotion. + # A burst paces; only deadline-aged queued waiters abandon the floor. + g = _incumbent_gen(0.0) + for uid in (7, 8, 9, 10): + _add_waiter(g, uid) + assert ar.resolve(g, 1.0) == ar.paced_ratio() + # burst promotes into the prompt batch: still paced + g._unprocessed_sequences.clear() + g._prompt_batch = FakeBatch([7, 8, 9, 10]) + assert ar.resolve(g, 2.0) == ar.paced_ratio() + + +def test_deadline_forces_stock(): + # A waiter queued behind a live paced train accrues pacing-attributable + # age (bounded 1 s per tick) and forces ratio 0 past the deadline. + g = _incumbent_gen(0.0) + _add_waiter(g) + assert ar.resolve(g, 1.0) == ar.paced_ratio() + g._prompt_batch = FakeBatch([9]) # a train is live; waiter 7 queued + t = 1.0 + while t < 11.0: # accrues 1.0 per 1 s tick + t += 1.0 + assert ar.resolve(g, t) == ar.paced_ratio() + assert ar.resolve(g, 12.0) == 0.0 # paced wait 11 > deadline 10 + # stood down is monotone: no further accrual, value holds past deadline + assert ar.resolve(g, 13.0) == 0.0 + + +def test_capacity_wait_does_not_age_toward_deadline(): + # No live train means the waiter is blocked by capacity (full decode + # batch), not pacing; ratio 0 could not admit it sooner, so it never + # ages and the floor holds indefinitely. + g = _incumbent_gen(0.0) + _add_waiter(g) + t = 0.0 + while t < 30.0: + t += 1.0 + assert ar.resolve(g, t) == ar.paced_ratio() + + +def test_pacing_survives_promotion_to_prompt_batch(): + # The regime B regression: the waiter's prompt is admitted into the + # prompt batch one tick after pacing starts. The chunk train is still + # competing prefill work; dropping pacing at promotion unpaces the + # whole prefill and puts the incumbent at the unpaced floor. + g = _incumbent_gen(0.0) + _add_waiter(g) + assert ar.resolve(g, 1.0) == ar.paced_ratio() + g._unprocessed_sequences.clear() + g._prompt_batch = FakeBatch([7]) + assert ar.resolve(g, 1.5) == ar.paced_ratio() + # prompt batch completes, rows reach decode: back to no waiters + g._prompt_batch = None + g._generation_batch = FakeBatch([101, 7]) + assert ar.resolve(g, 2.0) == 0.0 + + +def test_prefilling_waiter_does_not_age_toward_deadline(): + # A deep prompt's paced chunk train can far outlive deadline_s; its + # TTFT is bounded by the paced ratio itself ((1+r)x), so only queued + # waiters age toward the deadline. + g = _incumbent_gen(0.0) + _add_waiter(g) + ar.resolve(g, 1.0) + g._unprocessed_sequences.clear() + g._prompt_batch = FakeBatch([7]) + assert ar.resolve(g, 30.0) == ar.paced_ratio() + + +def test_queued_waiter_behind_prompt_batch_still_ages(): + g = _incumbent_gen(0.0) + _add_waiter(g) + ar.resolve(g, 1.0) + g._unprocessed_sequences.clear() + g._prompt_batch = FakeBatch([7]) # waiter 7 promoted; its train paced + _add_waiter(g, 8) # 8 queued behind it + t = 1.0 + while t < 11.0: + t += 1.0 + assert ar.resolve(g, t) == ar.paced_ratio() + assert ar.resolve(g, 13.0) == 0.0 # 8's paced wait > deadline 10 + + +def test_gate_hidden_ticks_freeze_deadline_and_keep_stamps(): + # The gate hides pending on declined ticks. Hidden ticks accrue no + # paced wait (the gate, not pacing, is why the waiter waits), and + # stamps survive so incumbency and accrued age resume, not reset. + g = _incumbent_gen(0.0) + _add_waiter(g) + g._prompt_batch = FakeBatch([9]) + for t in (1.0, 2.0, 3.0, 4.0, 5.0, 6.0): + assert ar.resolve(g, t) == ar.paced_ratio() + st = g._kq_auto + assert st.paced_wait[7] == 5.0 + seen = st.first_seen[7] + # gate declines: pending stashed empty, deferred dict marks uid 7 + stash = g._unprocessed_sequences + g._unprocessed_sequences = [] + g._kq_admit_deferred_s = {7: 0.0} + for t in (7.0, 8.0, 20.0): + ar.resolve(g, t) + assert st.paced_wait[7] == 5.0 # frozen, not accrued, not wiped + assert st.first_seen[7] == seen # arrival stamp survives hiding + # gate admits: pending restored; accrual resumes where it left off + g._unprocessed_sequences = stash + g._kq_admit_deferred_s = {} + for t in (21.0, 22.0, 23.0, 24.0, 25.0): + assert ar.resolve(g, t) == ar.paced_ratio() + assert ar.resolve(g, 27.0) == 0.0 # 5 + 6 accrued > deadline 10 + + +def test_suppressed_transition_logs_when_rate_window_opens(caplog): + # A transition inside the rate-limit window is deferred, not lost: + # the persisted state logs on a later tick once the window opens. + import logging + g = _incumbent_gen(0.0) + _add_waiter(g) + with caplog.at_level(logging.INFO, logger="gmlx.auto_ratio"): + ar.resolve(g, 1.0) # logs: pacing on + g._unprocessed_sequences.clear() + ar.resolve(g, 1.5) # off, suppressed (< 1 s) + assert "pacing off" not in caplog.text + ar.resolve(g, 2.1) # still off, window open: logs + assert "pacing off" in caplog.text + + +def test_log_rate_env_zero_logs_every_transition(caplog, monkeypatch): + import logging + monkeypatch.setenv("GMLX_DECODE_PREFILL_LOG_S", "0") + g = _incumbent_gen(0.0) + _add_waiter(g) + with caplog.at_level(logging.INFO, logger="gmlx.auto_ratio"): + ar.resolve(g, 1.0) + g._unprocessed_sequences.clear() + ar.resolve(g, 1.2) + assert "pacing on" in caplog.text and "pacing off" in caplog.text + + +def test_width_change_freezes_c_state(): + g = _incumbent_gen(0.0) + _add_waiter(g) + assert ar.resolve(g, 1.0) == ar.paced_ratio() + # admission widened the batch; width-2 bucket is only seeded, so the + # C state holds even though the chunk now looks cheap + g._generation_batch.uids.append(102) + g._kq_last_chunk_time = 0.01 + assert ar.resolve(g, 2.5) == ar.paced_ratio() + # real width-2 samples land: state re-evaluates and flips off + _feed_step(g, width=2) + assert ar.resolve(g, 4.0) == 0.0 + + +def test_kill_switch_resolves_static_paced(monkeypatch): + monkeypatch.setenv("GMLX_DECODE_PREFILL_AUTO", "0") + g = FakeGen() + assert ar.resolve(g, 0.0) == ar.paced_ratio() + + +def test_observe_skips_post_chunk_bracket(): + g = FakeGen() + ar.observe(g, 0.4, 0.4, 1) # chunk tick: poisons the next bracket + ar.observe(g, 9.9, 0.0, 1) # poisoned bracket: dropped + ar.observe(g, 0.05, 0.0, 1) # real step + st = g._kq_auto + assert st.s_by_width[1] == 0.05 + assert st.s_samples[1] == 1 + + +def test_floor_env_derives_both_constants(monkeypatch): + monkeypatch.setenv("GMLX_DECODE_PREFILL_FLOOR", "0.75") + assert ar.paced_ratio() == 3.0 + assert abs(ar.c_threshold() - 1.0 / 3.0) < 1e-9 diff --git a/tests/test_batch_sched.py b/tests/test_batch_sched.py index cb4c343..b8ac702 100644 --- a/tests/test_batch_sched.py +++ b/tests/test_batch_sched.py @@ -129,11 +129,41 @@ def test_ratio_flips_mid_run(paced, monkeypatch): assert g.chunks == paced_chunks + 10 # per-tick read: now stock +def test_auto_mode_wired_stock_without_incumbent(paced, monkeypatch): + monkeypatch.setenv("GMLX_DECODE_PREFILL_RATIO", "auto") + g = FakeGen() + for _ in range(10): + paced(g) + # burst shape (no admitted-row stamps precede the waiters): auto + # resolves 0 every tick => stock 1 decode : 1 chunk + assert g.chunks == 10 and g.decodes == 10 + + +def test_auto_kill_switch_paces_static(paced, monkeypatch): + monkeypatch.setenv("GMLX_DECODE_PREFILL_RATIO", "auto") + monkeypatch.setenv("GMLX_DECODE_PREFILL_AUTO", "0") + g = FakeGen() + for _ in range(40): + paced(g) + # auto disabled resolves to the static paced ratio (1.0 at rho 0.5) + assert g.decodes / g.chunks == 5.0 + + +def test_default_unset_is_auto(paced, monkeypatch): + monkeypatch.delenv("GMLX_DECODE_PREFILL_RATIO", raising=False) + g = FakeGen() + for _ in range(10): + paced(g) + # auto with no incumbent resolves 0: stock 1 decode : 1 chunk + assert g.chunks == 10 and g.decodes == 10 + + def test_bad_ratio_warns_and_defaults(paced, monkeypatch, caplog): monkeypatch.setenv("GMLX_DECODE_PREFILL_RATIO", "banana") g = FakeGen() with caplog.at_level("WARNING"): for _ in range(30): paced(g) - assert g.decodes / g.chunks > 2 # behaved as 1.0 + # falls back to the default (auto); burst shape resolves stock + assert g.chunks == 30 and g.decodes == 30 assert sum("banana" in r.message for r in caplog.records) == 1 diff --git a/tests/test_config.py b/tests/test_config.py index 5814ad0..11ab4a7 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -993,6 +993,10 @@ def test_decode_prefill_ratio_parsed_and_defaults_none(): ).decode_prefill_ratio == 1.5 # coerced to float assert build_config({"server": {"decode_prefill_ratio": 0}} ).decode_prefill_ratio == 0.0 # 0 => stock sched + assert build_config({"server": {"decode_prefill_ratio": "auto"}} + ).decode_prefill_ratio == "auto" + assert build_config({"server": {"decode_prefill_ratio": " AUTO "}} + ).decode_prefill_ratio == "auto" # case/space with pytest.raises(ConfigError): build_config({"server": {"decode_prefill_ratio": "fast"}})