diff --git a/CHANGELOG.md b/CHANGELOG.md index 5497bfe..abad8c9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,10 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### Added +- serve --speculative: a request arriving while one stream decodes with + MTP no longer waits for it to finish; the stream converts to shared + batch decode and speculation resumes once the batch drains back under + the width cap. GMLX_MTP_PREEMPT / GMLX_MTP_RESUME disable each half. - gmlx chat --server: chat against a running server as a plain client, without the assistant's tools and memory (no background requests). Engages automatically when the config's server is already up and diff --git a/docs/README.md b/docs/README.md index ceea470..0b1d083 100644 --- a/docs/README.md +++ b/docs/README.md @@ -112,6 +112,10 @@ script; do not edit it by hand. [serving-architecture.md](serving-architecture.md) explains how the pieces compose: loader, engine, batching, and the HTTP layers. +[speculative-batching.md](speculative-batching.md) covers how speculative +decoding and continuous batching run together: the two decode loops, the +width cap, and the preempt + resume transitions between them. + [adding-architectures.md](adding-architectures.md) is what adding a model family involves and the acceptance gate an architecture clears to be listed as supported. diff --git a/docs/performance.md b/docs/performance.md index ce49528..d658dc0 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -153,9 +153,11 @@ One interaction to know about: speculation and batching compete for the same bandwidth. Verifying a draft widens each request's weight reads, which is nearly free while one stream decodes and costly once several do, so the lift falls as concurrency rises. The server handles this for you with a per-model -batch-width cap: speculation runs while the live batch is narrow and the -batch finishes in plain decode once it grows past the cap, with the drafter -left loaded for the next one. +batch-width cap: speculation runs while the live batch is narrow, the batch +decodes plain past the cap, and speculation resumes once it drains back +under it. A lone speculating stream likewise yields to arriving requests +instead of making them wait. The transition mechanics are in +[speculative-batching.md](speculative-batching.md). Where the trade turns depends on the drafter and on whether the target routes experts. A native head verified by a dense hybrid-attention target keeps diff --git a/docs/server-config.md b/docs/server-config.md index a0d07b7..5f692ac 100644 --- a/docs/server-config.md +++ b/docs/server-config.md @@ -585,13 +585,16 @@ speculates only while at most N requests decode together. A drafter that can only handle one sequence clamps any larger value, since exceeding it raises rather than running slowly. -A batch that grows past the cap finishes in plain decode. The drafter stays -loaded and untouched, and the next batch to form re-evaluates. There is no -mid-flight switch back, so a continuously busy batch keeps decoding plain -until it drains. `GMLX_MTP_WIDTH_CAP` overrides every model at once (set it to -`0` to measure a model uncapped) and `--speculative-width-cap` does the same -from the CLI. The measured numbers behind the defaults are in -[performance.md](performance.md#mtp-speculative-decoding). +A batch that grows past the cap converts to plain decode with the drafter +left loaded, and once it drains back to the cap it re-arms and speculates +again (a capture round rebuilds the drafter state; mechanics in +[speculative-batching.md](speculative-batching.md)). `GMLX_MTP_WIDTH_CAP` +overrides every model at once (set it to `0` to measure a model uncapped) and +`--speculative-width-cap` does the same from the CLI. `GMLX_MTP_PREEMPT=0` +and `GMLX_MTP_RESUME=0` disable the batching transitions themselves (a lone +speculating stream then makes arriving requests wait, and a gated batch +stays plain until it finishes). The measured numbers behind the defaults are +in [performance.md](performance.md#mtp-speculative-decoding). An entry whose file is gone from disk does not stop the server: it is skipped with a log warning at startup (and on config reload), disappears from diff --git a/docs/speculative-batching.md b/docs/speculative-batching.md new file mode 100644 index 0000000..a3f6d74 --- /dev/null +++ b/docs/speculative-batching.md @@ -0,0 +1,146 @@ +# Speculative decoding under continuous batching + +How the server runs MTP speculation and continuous batching together: the +two decode loops, the batch-width cap, and the preempt + resume mechanics +that move requests between them without interrupting any token stream. + +For what speculation is and how to enable it, see +[performance.md](performance.md#mtp-speculative-decoding). For the width-cap +config key, see +[server-config.md](server-config.md#speculative_width_cap). + +## The two decode loops + +A speculative generation runs in one of two loops, chosen by live batch +width: + +- The scalar loop (one request decoding). The fastest path: draft and + target sampler RNG streams are kept coupled, which lets sampled drafts be + accepted against sampled targets and yields the highest acceptance rates. + This loop serves the common case of a single stream decoding at full + speculative speed. +- The batch loop (two or more requests). Tracks per-row state (bonus token, + KV offset, budget, finished flag), drafts greedily (coupled RNG does not + extend across rows), and checks a per-model width cap: a batch wider than + the cap decodes plain, because verification widens every row's weight + reads and past a measured knee the batch is faster without drafting. + +New requests join a running batch between verify rounds: the loop drains an +injection queue, extends the target KV cache and the drafter with the new +rows, and the width cap is re-checked against the widened batch. + +## Preempt: joining a scalar generation + +The scalar loop has no injection boundary; its speed comes from not being a +batch. Historically that meant a prefilled request arriving while a scalar +speculative generation streamed had to wait for the incumbent to finish +before starting its own decode. The wait is wrong on both axes: the waiter's +time to first token stretches to the incumbent's remaining generation, and +aggregate throughput loses too, because a single speculating stream is +slower than the same hardware decoding several streams plain. + +So the server preempts. When waiters queue against a live scalar +speculative generation: + +1. The scalar generator is closed at its verify-round boundary. Its cleanup + path rolls the target KV cache back to exactly the delivered tokens, so + the boundary state is clean by construction: the next undelivered token + (the round's bonus token) has no KV entry yet. +2. The generation is rebuilt as a batch-loop generator, restarting from that + bonus token with its real emitted count, but unarmed: no drafter state, + no captured hidden. Single-sequence caches are lifted to their batch + classes on the way. +3. The rebuilt loop's first injection drain admits the waiters. If the new + width exceeds the cap the batch decodes plain (the common case: any + second stream trips a cap of 1); otherwise the batch arms itself with a + capture round (below) and keeps speculating at the new width. + +The incumbent's stream continues without a gap. Its rate steps down from +solo-speculative to shared-plain while the batch is wide, which is the +correct trade: total tokens per second across streams goes up. + +`GMLX_MTP_PREEMPT=0` restores the old behavior (waiters hold until the +scalar generation drains). + +## Resume: re-arming a drained batch + +A batch gated to plain decode used to stay plain for the generator's life, +even after finishing rows brought it back under the cap. That latch existed +because re-arming a drafter mid-flight needs fresh hidden state and +shared-KV for every surviving row, and reusing stale per-row state was the +crash seam of an earlier campaign. + +The resume path re-arms without touching stale state, by re-running the +generator's own cold-start sequence on fresh captures: + +1. When a gated batch drains to the cap or below, the loop first finishes + consuming its plain-decode double buffer. Gated rounds dispatch the next + round's forward before reading this round's tokens, and that dispatched + step has already appended its KV; discarding it would corrupt the cache, + so one more plain round runs without dispatching a successor. +2. The next round is a capture round: a one-position verify forward of each + row's pending bonus token, with hidden-state and shared-KV capture on. + This emits one token per row at plain-decode cost. +3. The drafter is reset and cold-started from the capture: drafters that + teacher-force a prompt seed from target hidden accept the one-token + capture (draft quality ramps back over the next rounds), and shared-KV + drafters get their view re-set from the verify capture through the same + round tail every armed round uses. +4. Subsequent rounds speculate normally at the drained width. + +Rows within a small remaining-budget threshold are not worth the capture +cost and finish plain instead. A new admission landing in the same round +wins over a pending resume: the injection drain runs first and re-trips the +gate, so a batch never arms over the cap. + +`GMLX_MTP_RESUME=0` restores the one-way latch. + +## Semantics and caveats + +- Token streams are continuous across every transition. Preempt restarts + from the exact rollback boundary; resume consumes the plain lookahead + before capturing. Nothing is skipped, re-emitted, or re-sampled. +- A preempted request decodes under batch-loop semantics for the rest of + its generation, including after the batch drains back to a single row: + greedy drafting instead of the scalar loop's coupled sampling, which + costs a few points of acceptance at temperature. The next request starts + scalar again. +- A preempted request drops its prompt-cache retirement context: its prefix + is not offered back to the APC when it finishes. Waiters and later + requests retire normally. +- The capture round emits at plain-decode rate; the speculative speedup + returns on the round after. Resumes are therefore paced by the + remaining-budget threshold rather than fired for nearly-done rows. + +## Longer plays + +Two designs that would raise the width caps themselves rather than manage +around them. Documented here for a future pass; neither is built. + +### Ragged mixed verify forward + +Today every row in a verify round carries the same draft depth, so the +verify forward is a rectangle: batch width times block size. Rows with cold +drafters (fresh joins, fresh resumes) waste verify positions on drafts that +will not be accepted, and MoE targets pay the expert union of every +position in the rectangle. + +A ragged verify would give each row its own draft length, packing the +forward as one variable-length sequence batch (the runtime already has +ragged prefill machinery). The MoE win is the interesting one: expert +gather cost scales with the union of experts touched, so trimming wasted +positions trims real bandwidth, and the width-2 loss that currently caps +MoE targets at 1 was measured with rectangular verify. A ragged forward +re-opens that measurement. + +### Tree verify + +The caps encode a linear-draft trade: each drafted position must beat plain +decode for every row. Verifying a token tree per row instead of a chain +raises acceptance per verify forward (multiple continuations share a +prefix), which shifts the knee outward: the batched verify does more useful +work per unit of bandwidth, so speculation stays profitable at widths that +lose today. This changes the B > 1 verify arithmetic (attention masks over +tree positions, per-row acceptance walks over branches) and the drafter +contract (emit branching drafts), so it is a program, not a patch. The +scalar loop would gain too, but the batch knee is where the cap lives. diff --git a/gmlx/spec_engine.py b/gmlx/spec_engine.py index 614426b..7ce6f6d 100644 --- a/gmlx/spec_engine.py +++ b/gmlx/spec_engine.py @@ -21,7 +21,7 @@ import mlx.core as mx from . import prefill_decay -from .envflags import env_int +from .envflags import env_bool, env_int _log = logging.getLogger(__name__) @@ -1305,6 +1305,9 @@ def _buffered_extend(self, other): if not hasattr(self, "_pending_injections"): self._pending_injections = [] self._pending_injections.append(other) + _debug_note(f"[mtp] extend buffered: +{len(other._all_uids)} rows " + f"(pending={len(self._pending_injections)}, " + f"active={active})") SpecBatch.extend = _buffered_extend @@ -1338,19 +1341,74 @@ def _filter_with_release(self, keep): # 4. Process pending injections in next() before advancing the generator _orig_next = SpecBatch.next + def _note_last_tokens(self, responses) -> None: + # Last delivered token per uid: the bonus a preempt rebuild restarts + # from (its KV is not yet in the cache at a round boundary). + stash = getattr(self, "_kq_last_tokens", None) + if stash is None: + stash = self._kq_last_tokens = {} + for r in responses: + if r.token is not None: + stash[r.uid] = int(r.token) + + def _lift_host_cache(c): + """Promote a single-sequence host cache to its batch class so the + rebuilt batch generator can extend/filter it (same lift the + injection path applies to incoming caches).""" + if hasattr(c, "filter") and hasattr(c, "extend"): + return c + lifted = type(c).merge([c]) + stamp = getattr(c, "_gmlx_cascade", None) + if stamp is not None: + lifted._gmlx_cascade = stamp + return lifted + + 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. + + The rebuilt row carries no APC retirement context (batch-loop rows + start with retire_ctxs None), so the preempted request's prefix is + not offered back to the prompt cache when it finishes.""" + if not env_bool("GMLX_MTP_PREEMPT", True): + return False + if not getattr(self, "_sent_first", False): + return False + last = getattr(self, "_kq_last_tokens", {}).get(self._all_uids[0]) + if last is None: + return False + it = self._rounds_iter + if it is not None: + self._rounds_iter = None + it.close() + 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 + self.shared_kv_states = None + self.prompt_tokens = None + self.model._kq_rebuild_emitted = [int(self._num_tokens[0])] + _debug_note("[mtp] preempt: scalar generation rebuilt for " + "continuous batching") + return True + def _next_with_injection(self): pending = getattr(self, "_pending_injections", None) # Mid-flight adoption works only when the batch rounds generator is # running: it drains model._generator_injections at its round - # boundaries. The scalar (B=1) generator never does, so merging uids - # into a scalar batch strands the entry -- the injected request's - # continuation then re-dispatches from the wrong state (the finished - # row's cache) and its stream is silently truncated. Leave scalar - # injections buffered; _len_with_promotion adopts them wholesale - # (their own cache/hidden/first token) once the current request ends. + # boundaries. The scalar (B=1) generator never does, so a live + # scalar host is preempted first: its generator closes at the round + # boundary and the batch is rebuilt armless on the batch loop. # `_all_uids` is an mlx-vlm generator internal (stable under the # ==0.6.3 pin); re-verify this batch-vs-scalar signal on a pin lift. - if pending and len(self._all_uids) > 1: + preempted = False + if pending and len(self._all_uids) == 1: + preempted = _preempt_scalar(self) + if pending and (len(self._all_uids) > 1 or preempted): responses = [] gen_inj = getattr(self.model, "_generator_injections", None) if gen_inj is None: @@ -1396,10 +1454,12 @@ def _next_with_injection(self): more = _orig_next(self) responses.extend(more) + _note_last_tokens(self, responses) _release_if_finished(self) return responses responses = _orig_next(self) + _note_last_tokens(self, responses) _release_if_finished(self) return responses @@ -1453,7 +1513,10 @@ def _owned_server_rounds( ): batch_size = int(first_bonus.shape[0]) if first_bonus.ndim > 0 else 1 if draft_kind == "mtp": - if batch_size == 1: + # hidden=None marks a preempted scalar generation rebuilt for + # continuous batching: it must run the batch loop (arm-from- + # capture entry), never the scalar fast path. + if batch_size == 1 and hidden is not None: if not _first_use_b1[0]: _debug_note("[mtp] owned round: B=1 scalar path") _first_use_b1[0] = True diff --git a/gmlx/speculative.py b/gmlx/speculative.py index bff94aa..65b385f 100644 --- a/gmlx/speculative.py +++ b/gmlx/speculative.py @@ -1465,6 +1465,11 @@ def _lift_injected_cache(cache, other): _width_cap_memo: tuple[str, int | None] = ("", None) +# A gated batch back under the cap only re-arms when every surviving row +# still has at least this much budget left; re-arming for a nearly-done row +# costs more than it saves. +_RESUME_MIN_REMAINING = 32 + def _mtp_width_cap(drafter) -> int: """Effective MTP batch-width cap: speculate only while the live decode @@ -1590,9 +1595,9 @@ def _owned_decode_rounds_batch( lm, prompt_cache: list, *, - hidden: mx.array, + hidden: mx.array | None, b: list[int], - shared_kv: dict, + shared_kv: dict | None, seed_tokens: mx.array | None, emitted: list[int], max_tokens: int, @@ -1602,7 +1607,7 @@ def _owned_decode_rounds_batch( eos_token_ids: set | None = None, row_ids: list[int] | None = None, ) -> Iterator[tuple[list[int | None], Any]]: - """Owned batched MTP decode loop (B >= 2). + """Owned batched MTP decode loop (B >= 1 under continuous batching). Structurally parallel to the B=1 _owned_decode_rounds but tracks per-row state (bonus token, KV offset, emitted count, finished flag). Uses @@ -1617,6 +1622,13 @@ def _owned_decode_rounds_batch( New rows' target KV caches are extended into ``prompt_cache``, the drafter is prefilled in isolation then merged, and the per-row loop state grows to include the newcomers. + + ``hidden=None`` means arm-from-capture: no prefill hidden/shared-KV is + available (a preempted scalar generation rebuilt into this loop). The + first speculative round is then a capture round: an S=1 verify with + hidden/shared-KV capture on, whose state cold-starts the drafter. The + same capture round re-arms a width-gated batch that has drained back + under the cap (resume; GMLX_MTP_RESUME=0 disables). """ token_dtype = mx.int32 greedy = sampler is None @@ -1660,6 +1672,18 @@ def _owned_decode_rounds_batch( drafter._kq_head_request = None except Exception: pass # slotted/frozen drafter forbids ad-hoc attrs + def _reset_armed(n: int) -> None: + # B=1-only drafters raise on a left_padding list; fall back bare. + try: + drafter.reset(model, left_padding=[0] * n) + except (TypeError, ValueError): + 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 + # must be a capture round that arms the drafter. + need_arm = (not gated) and hidden is None + # reset() is bind + empty caches, no compute, so it runs gated too (that # keeps "the previous batch already released the drafter" off the critical # path). Without left_padding when gated: B=1-only drafters raise on a @@ -1667,7 +1691,7 @@ def _owned_decode_rounds_batch( if gated: drafter.reset(model) else: - drafter.reset(model, left_padding=[0] * len(b)) + _reset_armed(len(b)) sampler_rng = _SpeculativeSamplerRNG(drafter, enabled=False) draft_kwargs = {} @@ -1675,7 +1699,7 @@ def _owned_decode_rounds_batch( draft_kwargs["greedy"] = True draft_sampler = _argmax_sampler - if not gated: + if not gated and not need_arm: prefill_draft = getattr(drafter, "prefill_from_target_hidden", None) if callable(prefill_draft) and seed_tokens is not None: sampler_rng.draft_call( @@ -1691,14 +1715,14 @@ def _owned_decode_rounds_batch( # TypeError. Every remaining use of either is behind `not gated`. hidden = None shared_kv = None - else: + elif not need_arm: if hidden.shape[1] > 1: hidden = hidden[:, -1:, :] hidden = _mtp_draft_hidden(lm, hidden) L_prefill = _mtp_cache_offset_max(prompt_cache) positions = [L_prefill] * len(b) - if not gated: + if not gated and not need_arm: drafter.set_shared_kv( shared_kv, kv_offset=L_prefill, position=_mtp_draft_position(mx.array(positions)), @@ -1739,13 +1763,54 @@ def _gated_step(inputs): return sampler(logprobs).reshape(-1).astype(token_dtype) def _trip_width_gate(width: int) -> None: - """Latch the batch into plain decode for the rest of this generator.""" + """Latch the batch into plain decode until it drains under the cap.""" nonlocal gated gated = True _log_width_cap_once( f"trip: B={width} > cap={cap}; batch converts to plain decode " f"until drained") + def _resume_ready() -> bool: + """A gated batch back under the cap may re-arm and speculate again.""" + if not cap or len(active_idx) > cap: + return False + if not env_bool("GMLX_MTP_RESUME", True): + return False + if getattr(model, "_generator_injections", None): + return False + # Re-arming costs a capture forward + a drafter seed; skip it for + # rows about to finish anyway. + return all(max_tok[i] - emitted[i] >= _RESUME_MIN_REMAINING + for i in active_idx) + + def _arm_capture(): + """Capture round: an S=1 verify of the pending bonus tokens with + hidden/shared-KV capture on, then the drafter cold start from that + fresh state (the generator-entry sequence, never stale per-row + state). Emits one token per row through the shared round tail.""" + nonlocal hidden + b_arr = mx.array([b[i] for i in active_idx], dtype=token_dtype) + with mx.stream(generation_stream): + verify = _mtp_verify_target( + lm, b_arr[:, None], prompt_cache, sampler, + sample_target_tokens=greedy) + budgets = [max(1, max_tok[i] - emitted[i]) for i in active_idx] + accepted_list, new_tokens_list = _coupled_walk_batch( + lm, verify, mx.zeros((len(active_idx), 0), dtype=token_dtype), + _walk_sampler, budgets) + _reset_armed(len(active_idx)) + prefill_draft = getattr(drafter, "prefill_from_target_hidden", None) + if callable(prefill_draft): + # One (bonus, hidden) pair per row; the seed contract accepts a + # 1-token capture (draft quality ramps back up over rounds). + next_b = mx.array([nt[-1] for nt in new_tokens_list], + dtype=token_dtype) + sampler_rng.draft_call( + prefill_draft, b_arr[:, None], verify.hidden, next_b, + draft_sampler, token_dtype, **draft_kwargs) + hidden = _mtp_draft_hidden(lm, verify.hidden[:, -1:, :]) + return accepted_list, new_tokens_list, verify + def _drain_injections(): # continuous-batch injection nonlocal hidden, B_orig, _gated_pending @@ -1772,7 +1837,10 @@ def _drain_injections(): cache, inj["prompt_cache"][i]) extend_fn(other) - if not gated: + # hidden is None while un-armed (arm-from-capture entry): + # the drafter holds no state to inject into, and the capture + # round will build hidden/shared-KV for every row at once. + if not gated and hidden is not None: inject_fn = getattr(drafter, "inject_rows", None) if callable(inject_fn): inject_fn( @@ -1780,7 +1848,6 @@ def _drain_injections(): inj["first_tokens"], draft_sampler, token_dtype, greedy=True) - if not gated: inj_hidden = inj["hidden"] if inj_hidden.shape[1] > 1: inj_hidden = inj_hidden[:, -1:, :] @@ -1814,7 +1881,7 @@ def _drain_injections(): max_tok.append(int(inj_max[row])) B_orig += 1 - if _needs_shared_kv and not gated: + if _needs_shared_kv and not gated and shared_kv is not None: # The old raw batch-axis concat crashed on ragged seq # and, when set_shared_kv had normalized into a copy, # never reached the drafter anyway -- whose stored @@ -1856,6 +1923,21 @@ def _drain_injections(): # ratio), so re-check even without an admission. if not gated and cap and len(active_idx) > cap: _trip_width_gate(len(active_idx)) + # Resume: a gated batch drained back under the cap re-enters + # speculation via a capture round. The gated double buffer holds a + # 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 + if gated and _resume_ready(): + if _gated_pending is None: + gated = False + need_arm = True + _log_width_cap_once( + f"resume: B={len(active_idx)} <= cap={cap}; " + f"re-arming drafter") + else: + dispatch_next = False n_active = len(active_idx) if gated: @@ -1883,8 +1965,11 @@ def _drain_injections(): # this leaves in the cache is past every retirement's store_len # (retirement is driven by `positions`, not the cache offset) and # nothing reads the cache offset while gated. - _gated_pending = _gated_step(toks) - mx.async_eval(_gated_pending) + if dispatch_next: + _gated_pending = _gated_step(toks) + mx.async_eval(_gated_pending) + else: + _gated_pending = None # Under GMLX_ROUND_PROFILE the columns read as dispatch (CPU build # of round N+1) / wait (GPU finishing round N) / emit-prep instead # of draft / verify / walk. @@ -1895,6 +1980,15 @@ def _drain_injections(): accepted_list = [0] * n_active new_tokens_list = [[int(t)] for t in toks.tolist()] _t1 = time.perf_counter() + elif need_arm: + need_arm = False + _t0 = time.perf_counter() + _gap = (_t0 - _prev_end) * 1e3 if _prev_end else 0.0 + bs = 1 + max_a = 0 + accepted_list, new_tokens_list, verify = _arm_capture() + _td = _tv = time.perf_counter() if _ROUND_PROFILE else _t0 + _t1 = time.perf_counter() else: remaining = [ max(1, max_tok[active_idx[j]] - emitted[active_idx[j]] + 1) @@ -1993,6 +2087,10 @@ def _drain_injections(): rejected_global = bs - (max_a + 1) next_shared_kv = _slice_shared_kv_batch( verify.shared_kv_states, rejected_global, accepted_list, max_a) + # Track the freshest slice so an injection merge has a live + # fallback dict (an armed-from-capture generator starts with + # shared_kv=None; the prefill-armed one would go stale). + shared_kv = next_shared_kv else: next_shared_kv = shared_kv @@ -2148,11 +2246,22 @@ def owned_server_rounds_batch( b = first_bonus.reshape(-1).tolist() _buffer_mtp_target_cache(prompt_cache, drafter, draft_block_size) eff_sampler = None if greedy_sampling else sampler + # A preempted scalar generation rebuilt into this loop carries its real + # emitted counts (the SpecBatch stop_check owns the hard budget; this + # only keeps block sizing honest near the end). + emitted = getattr(model, "_kq_rebuild_emitted", None) + if emitted is not None: + try: + del model._kq_rebuild_emitted + except AttributeError: + pass + if not emitted or len(emitted) != B: + emitted = [1] * B yield from _owned_decode_rounds_batch( model, drafter, lm, prompt_cache, hidden=hidden, b=b, shared_kv=shared_kv_states, seed_tokens=prompt_tokens, - emitted=[1] * B, max_tokens=max_tokens, + emitted=list(emitted), max_tokens=max_tokens, sampler=eff_sampler, draft_block_size=draft_block_size, stop_check=stop_check, eos_token_ids=eos_token_ids, row_ids=row_ids) diff --git a/tests/test_mtp_preempt_resume.py b/tests/test_mtp_preempt_resume.py new file mode 100644 index 0000000..b2d99ea --- /dev/null +++ b/tests/test_mtp_preempt_resume.py @@ -0,0 +1,482 @@ +"""MTP preempt + resume: a live B=1 speculative generation preempts at a +verify-round boundary when prefilled waiters queue (the scalar generator +closes, the batch rebuilds armless on the batch loop, waiters join through +the injection drain), and a width-gated batch that drains back under the cap +resumes speculation through a capture round. + +Loop-level tests drive _owned_decode_rounds_batch with an armable fake +drafter whose draft_block returns an EMPTY draft, so armed rounds run the +full speculative machinery (S=1 verify + zero-draft walk) with deterministic +echo streams. Engine-level tests run the real SpeculativeGenerationBatch +against a recording fake rounds generator, the same pattern as +test_spec_engine_release. +""" + +from types import SimpleNamespace + +import mlx.core as mx +import pytest + +from gmlx.speculative import _owned_decode_rounds_batch, _width_cap_logged +import gmlx.speculative as spec + +from test_mtp_width_cap import _FakeCache, _StrictDrafter + +VOCAB = 32 + + +@pytest.fixture(autouse=True) +def _clean_env(monkeypatch): + for var in ("GMLX_MTP_WIDTH_CAP", "MLX_VLM_GGUF_SPEC_WIDTH_CAP", + "GMLX_MTP_PREEMPT", "GMLX_MTP_RESUME"): + monkeypatch.delenv(var, raising=False) + spec._width_cap_memo = ("", None) + _width_cap_logged.clear() + yield + spec._width_cap_memo = ("", None) + _width_cap_logged.clear() + + +class _VerifyEchoLM: + """Echo target (next token = input + 1) that also serves the + plain-forward verify branch (hidden_states / shared_kv / logits), so + capture rounds and zero-draft speculative rounds run end to end. + Plain and verify forwards are recorded separately: the consume-then- + capture protocol is asserted from these counters.""" + + def __init__(self): + self.plain_widths = [] + self.verify_widths = [] + self._rope_deltas = None + + def __call__(self, x, cache=None, return_hidden=False, + return_shared_kv=False, **kw): + B, S = x.shape + nxt = (x + 1) % VOCAB + onehot = (mx.arange(VOCAB)[None, None, :] == nxt[:, :, None]) + logits = onehot.astype(mx.float32) * 10.0 + if return_hidden: + self.verify_widths.append(B) + return SimpleNamespace( + logits=logits, + hidden_states=[mx.zeros((B, S, 8))], + shared_kv_states={"full": (mx.zeros((B, 2, S, 4)), + mx.zeros((B, 2, S, 4)))}, + gdn_states=None) + self.plain_widths.append(B) + return SimpleNamespace(logits=logits) + + +class _ArmableDrafter: + """Cold-startable fake: the seed calls (reset / prefill_from_target_hidden + / set_shared_kv) are recorded, draft_block returns an empty draft so armed + rounds exercise verify + walk without draft-quality modeling.""" + + uses_shared_kv = True + + def __init__(self, cap=2, block_size=4): + self.config = SimpleNamespace(block_size=block_size) + self.accept_lens = [] + self.draft_lens = [] + self.mtp_width_cap = cap + self.mtp_width_limit = 0 + self.reset_calls = [] + self.prefill_calls = [] + self.shared_kv_calls = [] + self.draft_calls = [] + + def reset(self, model, left_padding=None): + self.reset_calls.append(left_padding) + return [] + + def set_shared_kv(self, shared_kv, kv_offset=0, position=None, + kv_valid_len=None, left_padding=None): + self.shared_kv_calls.append(kv_offset) + + def prefill_from_target_hidden(self, tokens, hidden, next_tokens, + sampler, dtype, **kw): + self.prefill_calls.append(tuple(hidden.shape)) + + def draft_block(self, b, hidden, kv, n, sampler, dtype, **kw): + self.draft_calls.append(int(b.shape[0])) + return mx.zeros((int(b.shape[0]), 0), dtype=dtype) + + +def _drive_armless(drafter, *, B, max_tokens, lm=None, model=None, + rounds=None, stop_check=None): + """Run the batch loop from an armless start (hidden=None, shared_kv=None), + the state a preempted scalar generation rebuilds into.""" + model = model if model is not None else SimpleNamespace() + lm = lm if lm is not None else _VerifyEchoLM() + prompt_cache = [_FakeCache(width=B)] + gen = _owned_decode_rounds_batch( + model, drafter, lm, prompt_cache, + hidden=None, + b=list(range(1, B + 1)), + shared_kv=None, + seed_tokens=None, + emitted=[1] * B, + max_tokens=max_tokens, + sampler=None, + draft_block_size=None, + stop_check=stop_check, + ) + out = [] + for toks, meta in gen: + out.append((toks, meta)) + if rounds is not None and len(out) >= rounds: + gen.close() + break + return out, lm, prompt_cache + + +# -- armless entry (rebuilt after preempt) -------------------------------- + + +def test_armless_entry_over_cap_gates_plain(): + """A rebuilt batch born over the cap decodes plain with hidden=None + everywhere: nothing on the gated path may dereference it.""" + d = _StrictDrafter(cap=2) + lm = _VerifyEchoLM() + out, _, _ = _drive_armless(d, B=3, max_tokens=4, lm=lm) + assert d.forward_calls == [] + assert d.shared_kv_calls == 0 + assert lm.verify_widths == [] + # echo chains from seeds [1, 2, 3] + assert [toks for toks, _ in out] == [[2, 3, 4], [3, 4, 5], [4, 5, 6]] + + +def test_armless_entry_capture_arms_then_speculates(): + """Under the cap, the first round is a capture round (S=1 verify, drafter + cold start) and later rounds run the speculative machinery. The token + stream is the same plain echo chain: arming must not skip or duplicate.""" + d = _ArmableDrafter(cap=0) + out, lm, _ = _drive_armless(d, B=2, max_tokens=4) + assert [toks for toks, _ in out] == [[2, 3], [3, 4], [4, 5]] + # every round is a verify (capture round + zero-draft rounds), no plain + assert lm.verify_widths == [2, 2, 2] + assert lm.plain_widths == [] + # cold start ran exactly once, from the 1-token capture + assert d.prefill_calls == [(2, 1, 8)] + # rounds after the capture actually drafted (empty blocks) + assert d.draft_calls == [2, 2] + # tail re-set the drafter's shared-KV view each armed round except the + # last (the all-finished break exits before the tail) + assert len(d.shared_kv_calls) == 2 + + +def test_armless_capture_skips_entry_seed(): + """The entry seed block (prefill from seed_tokens + set_shared_kv at + L_prefill) reads state an armless start does not have; the capture round + is the only seed path.""" + d = _ArmableDrafter(cap=0) + _drive_armless(d, B=2, max_tokens=2) + assert len(d.prefill_calls) == 1 + + +# -- resume (gated batch drains under the cap) ---------------------------- + + +def _finish_row0(orig, tok): + return orig == 0 + + +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 + is already in the cache), the round after arms via capture, and the + stream stays the exact echo chain across both transitions.""" + d = _ArmableDrafter(cap=2) + out, lm, _ = _drive_armless( + d, B=3, max_tokens=40, rounds=5, stop_check=_finish_row0) + assert [toks for toks, _ in out] == [ + [2, 3, 4], # gated round; row 0 finishes + [None, 4, 5], # consume round: dispatched lookahead, no re-prime + [None, 5, 6], # capture round (arm) + [None, 6, 7], # speculative round + [None, 7, 8], + ] + # gated rounds: prime + one dispatch at width 3, then nothing plain -- + # a discarded (re-primed) buffer would show a third plain forward + assert lm.plain_widths == [3, 3] + # capture + speculative rounds verify at the drained width + assert lm.verify_widths == [2, 2, 2] + assert d.prefill_calls == [(2, 1, 8)] + assert d.draft_calls == [2, 2] + + +def test_resume_env_kill_switch(monkeypatch): + """GMLX_MTP_RESUME=0 keeps the drained batch on plain decode (the old + latch behavior); the strict drafter would raise on any arm attempt.""" + monkeypatch.setenv("GMLX_MTP_RESUME", "0") + d = _StrictDrafter(cap=2) + out, lm, _ = _drive_armless( + d, B=3, max_tokens=40, rounds=5, stop_check=_finish_row0, + lm=_VerifyEchoLM()) + assert d.forward_calls == [] + assert len(out) == 5 + assert out[-1][0] == [None, 7, 8] + + +def test_resume_skipped_near_budget_end(): + """Re-arming costs a capture forward + a drafter seed; rows within + _RESUME_MIN_REMAINING of their budget finish plain instead.""" + d = _StrictDrafter(cap=2) + out, _, _ = _drive_armless( + d, B=3, max_tokens=8, lm=_VerifyEchoLM(), stop_check=_finish_row0) + assert d.forward_calls == [] + # rows 1..2 run to their budget (emitted 1 -> 8 = 7 rounds) + assert len(out) == 7 + + +def test_resume_rechecks_cap_on_new_admission(): + """An injection landing in the same round a resume would fire must win: + the drain runs first and re-trips the gate, so the batch never arms over + the cap.""" + d = _StrictDrafter(cap=2) + model = SimpleNamespace() + lm = _VerifyEchoLM() + prompt_cache = [_FakeCache(width=3)] + gen = _owned_decode_rounds_batch( + model, d, lm, prompt_cache, + hidden=None, b=[1, 2, 3], shared_kv=None, seed_tokens=None, + emitted=[1, 1, 1], max_tokens=40, sampler=None, + draft_block_size=None, stop_check=_finish_row0) + next(gen) # gated round; row 0 finishes -> batch at the cap + next(gen) # consume round; a resume is now due next round + model._generator_injections = [{ + "uids": ["w", "x"], + "prompt_cache": [_FakeCache(width=2, offset=9)], + "hidden": mx.zeros((2, 1, 8)), + "prompt_tokens": mx.zeros((2, 4), dtype=mx.int32), + "first_tokens": mx.array([7, 8], dtype=mx.int32), + "first_tokens_list": [7, 8], + "shared_kv_states": None, + }] + toks, _ = next(gen) + gen.close() + assert d.forward_calls == [] # never armed + assert len(toks) == 5 # both admissions joined the round + + +# -- rebuilt-generator emitted override ----------------------------------- + + +def test_rebuild_emitted_override_consumed(): + """owned_server_rounds_batch picks up model._kq_rebuild_emitted (the + preempted host's real emitted count), applies it to the budget, and + deletes the attribute.""" + from gmlx.speculative import owned_server_rounds_batch + + d = _ArmableDrafter(cap=0) + lm = _VerifyEchoLM() + model = SimpleNamespace() + model._kq_rebuild_emitted = [5] + gen = owned_server_rounds_batch( + model, d, [_FakeCache(width=1)], + None, + first_bonus=mx.array([3], dtype=mx.int32), + max_tokens=8, + sampler=None, + shared_kv_states=None, + prompt_tokens=None, + greedy_sampling=True, + ) + # patch lm resolution: owned_server_rounds_batch derives lm from model + model.language_model = lm + out = [toks for toks, _ in gen] + assert not hasattr(model, "_kq_rebuild_emitted") + # emitted resumes at 5, so budget 8 leaves exactly 3 more tokens + assert out == [[4], [5], [6]] + + +def test_rebuild_emitted_length_mismatch_ignored(): + """A stale or mis-sized override must not survive into an unrelated + batch: wrong length falls back to the fresh-generator default.""" + from gmlx.speculative import owned_server_rounds_batch + + d = _ArmableDrafter(cap=0) + lm = _VerifyEchoLM() + model = SimpleNamespace() + model._kq_rebuild_emitted = [5] + gen = owned_server_rounds_batch( + model, d, [_FakeCache(width=2)], + None, + first_bonus=mx.array([3, 4], dtype=mx.int32), + max_tokens=3, + sampler=None, + shared_kv_states=None, + prompt_tokens=None, + greedy_sampling=True, + ) + model.language_model = lm + out = [toks for toks, _ in gen] + assert not hasattr(model, "_kq_rebuild_emitted") + assert out == [[4, 5], [5, 6]] + + +# -- engine-level preempt (real SpeculativeGenerationBatch) --------------- + + +class _HostCache: + """Single-sequence cache: no filter/extend, so the preempt must lift it + through type(c).merge([c]).""" + + def __init__(self): + self._gmlx_cascade = "stamp" + + @classmethod + def merge(cls, caches): + lifted = _BatchCache() + lifted.merged_from = list(caches) + return lifted + + +class _BatchCache: + merged_from = None + + def filter(self, keep): + pass + + def extend(self, other): + pass + + +class _EngineDrafter: + def __init__(self): + self.reset_calls = 0 + + def reset(self, model, left_padding=None): + self.reset_calls += 1 + return [] + + +def _make_batch(ar, *, uids=(0,), model=None, max_tokens=6, cache=None): + return ar.SpeculativeGenerationBatch( + model=model if model is not None else SimpleNamespace(), + draft_model=_EngineDrafter(), + draft_kind="mtp", + uids=list(uids), + first_tokens=mx.array([5 + u for u in uids]), + prompt_cache=[cache if cache is not None else _HostCache()], + sampler=None, + stop_criteria=lambda tok: False, + max_tokens=[max_tokens] * len(uids), + hidden=mx.zeros((len(uids), 4, 8)), + shared_kv_states={"full": None}, + prompt_tokens=mx.array([[1, 2, 3]] * len(uids)), + greedy_sampling=True, + ) + + +def _recording_rounds(calls): + """Fake rounds generator: records its (hidden, first_bonus, cache) per + call, drains model._generator_injections the way the real batch loop + does (widening its yields), and flags close.""" + + def fake_rounds(model, draft_model, prompt_cache, hidden, **kw): + entry = {"hidden": hidden, "first_bonus": kw.get("first_bonus"), + "cache": list(prompt_cache), "closed": False} + calls.append(entry) + width = int(entry["first_bonus"].shape[0]) + base = 100 * len(calls) + n = 0 + try: + while True: + inj = getattr(model, "_generator_injections", None) + if inj: + width += sum(len(e["uids"]) for e in inj) + inj.clear() + n += 1 + yield [base + n] * width, None + finally: + entry["closed"] = True + + return fake_rounds + + +def test_preempt_rebuilds_scalar_for_waiters(monkeypatch): + 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", + _recording_rounds(calls)) + + model = SimpleNamespace() + host = _make_batch(ar, uids=(0,), model=model) + assert [r.token for r in host.next()] == [5] # first token + assert [r.token for r in host.next()] == [101] # scalar round + + waiter = _make_batch(ar, uids=(7,), model=model) + host.extend(waiter) # buffered + + responses = host.next() + # scalar generator closed at its round boundary + assert calls[0]["closed"] is True + # rebuilt armless from the last delivered token + assert len(calls) == 2 + assert calls[1]["hidden"] is None + assert calls[1]["first_bonus"].tolist() == [101] + # host cache lifted to the batch class, cascade stamp preserved + lifted = calls[1]["cache"][0] + assert isinstance(lifted, _BatchCache) + assert isinstance(lifted.merged_from[0], _HostCache) + assert lifted._gmlx_cascade == "stamp" + # real emitted count handed to the rebuilt generator + assert model._kq_rebuild_emitted == [2] + # waiter delivered its first token and joined the same round + assert [(r.uid, r.token) for r in responses] == [ + (7, 12), (0, 201), (7, 201)] + + +def test_preempt_env_kill_switch(monkeypatch): + from mlx_vlm.generate import ar + from gmlx.spec_engine import install_continuous_batch_admission + + install_continuous_batch_admission() + monkeypatch.setenv("GMLX_MTP_PREEMPT", "0") + calls = [] + monkeypatch.setattr(ar, "run_speculative_server_rounds", + _recording_rounds(calls)) + + model = SimpleNamespace() + host = _make_batch(ar, uids=(0,), model=model) + host.next() + host.next() + host.extend(_make_batch(ar, uids=(7,), model=model)) + + responses = host.next() + # old behavior: scalar keeps running, waiter stays buffered until drain + assert len(calls) == 1 + assert calls[0]["closed"] is False + assert [(r.uid, r.token) for r in responses] == [(0, 102)] + assert len(host._pending_injections) == 1 + + +def test_preempt_waits_for_first_delivery(monkeypatch): + """A batch that has not delivered its first tokens has no bonus to + rebuild from; the preempt fires on the following next() instead.""" + 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", + _recording_rounds(calls)) + + model = SimpleNamespace() + host = _make_batch(ar, uids=(0,), model=model) + host.extend(_make_batch(ar, uids=(7,), model=model)) + + assert [r.token for r in host.next()] == [5] # no preempt yet + assert calls == [] + + responses = host.next() + assert len(calls) == 1 + assert calls[0]["hidden"] is None + assert calls[0]["first_bonus"].tolist() == [5] + assert model._kq_rebuild_emitted == [1] + assert [(r.uid, r.token) for r in responses] == [ + (7, 12), (0, 101), (7, 101)]