feat(mtp): preempt and resume speculation under continuous batching

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

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.

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

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

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.

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