qwen3_5: load in-checkpoint MTP head + speculative rollback for hybrid (GDN) caches

mlx_lm/generate.py

@@ -526,8 +526,17 @@ def speculative_generate_step(
         model_cache = prompt_cache[: len(model.layers)]
         draft_cache = prompt_cache[len(model.layers) :]
 
-    if not cache.can_trim_prompt_cache(model_cache):
-        types = {type(c).__name__ for c in model_cache if not c.is_trimmable()}
+    def _can_speculate(c):
+        # Trimmable directly, or able to record an exact rollback while
+        # speculating (see ArraysCache.record_rollback). The model must also
+        # declare support: recording is done by its recurrent layers.
+        return c.is_trimmable() or (
+            hasattr(c, "record_rollback")
+            and getattr(model, "supports_speculative_rollback", False)
+        )
+
+    if not all(_can_speculate(c) for c in model_cache):
+        types = {type(c).__name__ for c in model_cache if not _can_speculate(c)}
         raise ValueError(
             f"Speculative decoding requires a trimmable prompt cache " f"(got {types})."
         )
@@ -604,13 +613,26 @@ def _draft_generate(y, num_draft):
         draft_y = _prefill(draft_model, draft_cache, y)
         y = _prefill(model, model_cache, y)
 
+    # After prefill (recording a rollback for the whole prompt would hold on to
+    # prompt-sized tensors), let rollback-capable caches start recording so
+    # verify steps can be trimmed exactly (no-op for regular KV caches). The
+    # draft cache records too: hybrid draft models (e.g. a small Qwen3.5) need
+    # their recurrent state rewound just like the target's.
+    for c in model_cache + draft_cache:
+        c.start_speculation()
+
     ntoks = 0
     # Set these so the finally block doesn't raise
     num_draft = 0
     n = 0
     try:
         while True:
             num_draft = min(max_tokens - ntoks, num_draft_tokens)
+            # Until this round's verify step has advanced the caches there is
+            # nothing to rewind; keep n == num_draft so an exception here (or a
+            # generator close) makes the finally-block rewind a no-op instead
+            # of trimming with values from a previous round.
+            n = num_draft
             draft_tokens = _draft_generate(draft_y, num_draft)
             if prev_tokens is not None:
                 prev_tokens = prev_tokens[: prev_tokens.size - y.size - num_draft + 1]
@@ -652,6 +674,8 @@ def _draft_generate(y, num_draft):
             _rewind_cache(num_draft, n)
     finally:
         _rewind_cache(num_draft, n)
+        for c in model_cache:
+            c.stop_speculation()
 
 
 def stream_generate(

mlx_lm/models/cache.py

@@ -146,6 +146,19 @@ def meta_state(self, v):
     def is_trimmable(self):
         return False
 
+    def start_speculation(self):
+        """Called by speculative decoding before draft/verify steps begin.
+
+        Caches that cannot trim directly but support an exact rollback (e.g.
+        recurrent-state caches, see ``ArraysCache``) use this to start recording
+        the information needed to roll a verify step back. No-op by default.
+        """
+        pass
+
+    def stop_speculation(self):
+        """Called by speculative decoding when it finishes. No-op by default."""
+        pass
+
     def size(self):
         """
         Return the size (i.e. sequence length) of the cache.
@@ -409,6 +422,17 @@ def nbytes(self):
 
 class RotatingKVCache(_BaseCache):
     step = 256
+    # Tokens of exact-rollback history kept while speculating (mirrors
+    # ArraysCache; a speculative trim never exceeds the last verify window).
+    _ROLLBACK_WINDOW = 64
+
+    def __new__(cls, *args, **kwargs):
+        # Set on __new__ (not __init__) so caches reconstructed by
+        # load_prompt_cache (which bypasses __init__) still have them.
+        instance = super().__new__(cls)
+        instance.speculating = False
+        instance._rollbacks = deque()
+        return instance
 
     def __init__(self, max_size, keep=0):
         self.keep = keep
@@ -418,6 +442,32 @@ def __init__(self, max_size, keep=0):
         self.max_size = max_size
         self._idx = 0
 
+    def start_speculation(self):
+        self.speculating = True
+        self._rollbacks.clear()
+
+    def stop_speculation(self):
+        self.speculating = False
+        self._rollbacks.clear()
+
+    def record_rollback(self, num_tokens, snapshot, keys, values):
+        """Recorded by ``update_and_fetch`` during a forward while
+        ``speculating``. A sliding-window KV cache is not directly trimmable
+        once it has wrapped (``offset >= max_size``): the tokens a verify step
+        pushes out of the window are gone. So we stash the pre-forward window
+        (``snapshot``) plus this forward's new K/V; ``trim`` rebuilds the exact
+        window for any accepted prefix by restoring the snapshot and
+        re-appending the first ``m`` tokens. The verify path uses
+        ``_update_concat`` which reallocates (never mutates in place), so
+        holding references in ``snapshot`` is safe."""
+        self._rollbacks.append((num_tokens, snapshot, keys, values))
+        total = sum(r[0] for r in self._rollbacks)
+        while (
+            len(self._rollbacks) > 1
+            and total - self._rollbacks[0][0] >= self._ROLLBACK_WINDOW
+        ):
+            total -= self._rollbacks.popleft()[0]
+
     def _trim(self, trim_size, v, append=None):
         to_cat = []
         if trim_size > 0:
@@ -510,6 +560,13 @@ def _update_in_place(self, keys, values):
         return self.keys, self.values
 
     def update_and_fetch(self, keys, values):
+        if self.speculating:
+            self.record_rollback(
+                keys.shape[2],
+                [self.keys, self.values, self._idx, self.offset],
+                keys,
+                values,
+            )
         if keys.shape[2] == 1:
             return self._update_in_place(keys, values)
         return self._update_concat(keys, values)
@@ -540,12 +597,43 @@ def meta_state(self, v):
         )
 
     def is_trimmable(self):
-        return self.offset < self.max_size
+        # While speculating every forward records an exact rollback, so the
+        # cache is trimmable within the recorded window even after it wraps.
+        return self.speculating or self.offset < self.max_size
 
     def trim(self, n):
-        n = min(self.offset, n)
-        self.offset -= n
-        self._idx -= n
+        if not self.speculating:
+            n = min(self.offset, n)
+            self.offset -= n
+            self._idx -= n
+            return n
+        recorded = sum(r[0] for r in self._rollbacks)
+        if recorded < n:
+            raise RuntimeError(
+                f"Cannot trim {n} tokens from RotatingKVCache: only {recorded} "
+                "tokens of exact rollback are recorded (speculative window)."
+            )
+        trimmed = 0
+        while trimmed < n:
+            num_tokens, snap, keys, values = self._rollbacks.pop()
+            take = min(n - trimmed, num_tokens)
+            m = num_tokens - take
+            # Restore the pre-forward window, then re-append the first m tokens
+            # of this chunk to reconstruct the exact accepted-prefix window.
+            self.keys, self.values, self._idx, self.offset = (
+                snap[0],
+                snap[1],
+                snap[2],
+                snap[3],
+            )
+            if m > 0:
+                if m == 1:
+                    self._update_in_place(keys[..., :1, :], values[..., :1, :])
+                else:
+                    self._update_concat(keys[..., :m, :], values[..., :m, :])
+                # The record still describes its first m tokens; keep it.
+                self._rollbacks.append((m, snap, keys, values))
+            trimmed += take
         return n
 
     def to_quantized(self, group_size: int = 64, bits: int = 4) -> QuantizedKVCache:
@@ -592,17 +680,81 @@ def nbytes(self):
 
 
 class ArraysCache(_BaseCache):
+    # Tokens of exact-rollback history kept while speculating — records beyond
+    # this window are dropped (speculative trims never exceed the last verify
+    # window, so 64 tokens is far more than any round needs). Bounding by
+    # tokens keeps the retained per-token stashes small.
+    _ROLLBACK_WINDOW = 64
+
     def __new__(cls, *args, **kwargs):
         instance = super().__new__(cls)
         instance.left_padding = None
         instance.lengths = None
+        instance.speculating = False
+        instance._rollbacks = deque()
         return instance
 
     def __init__(self, size, left_padding: Optional[List[int]] = None):
         self.cache = [None] * size
         if left_padding:
             self.left_padding = mx.array(left_padding)
 
+    def start_speculation(self):
+        self.speculating = True
+        self._rollbacks.clear()
+
+    def stop_speculation(self):
+        self.speculating = False
+        self._rollbacks.clear()
+
+    def record_rollback(self, num_tokens, fn, snapshot):
+        """Recorded by the owning layer during a forward while ``speculating``.
+
+        Args:
+            num_tokens (int): Sequence length of the forward being recorded.
+            fn (callable): ``fn(m) -> list`` returning the cache entries as they
+                would be had only the first ``m`` of ``num_tokens`` tokens been
+                processed. Must be exact (e.g. replay the recurrence from the
+                pre-forward state over the stashed per-token inputs).
+            snapshot (list): The cache entries from before the forward (returned
+                for a full rollback, ``m == 0``).
+        """
+        self._rollbacks.append((num_tokens, fn, snapshot))
+        total = sum(r[0] for r in self._rollbacks)
+        while (
+            len(self._rollbacks) > 1
+            and total - self._rollbacks[0][0] >= self._ROLLBACK_WINDOW
+        ):
+            total -= self._rollbacks.popleft()[0]
+
+    def is_trimmable(self):
+        # While speculating, every forward records an exact rollback, so the
+        # cache is trimmable within the recorded window.
+        return self.speculating
+
+    def trim(self, n):
+        if n <= 0:
+            return 0
+        if sum(r[0] for r in self._rollbacks) < n:
+            raise RuntimeError(
+                f"Cannot trim {n} tokens from ArraysCache: only "
+                f"{sum(r[0] for r in self._rollbacks)} tokens of exact rollback "
+                "are recorded. Recurrent state cannot be trimmed beyond the "
+                "speculative window."
+            )
+        trimmed = 0
+        while trimmed < n:
+            num_tokens, fn, snapshot = self._rollbacks.pop()
+            take = min(n - trimmed, num_tokens)
+            m = num_tokens - take
+            self.cache = list(snapshot) if m == 0 else fn(m)
+            trimmed += take
+            if m > 0:
+                # The record still describes the first m tokens of its chunk
+                # (fn(m') is valid for any m' <= m), so keep it for further trims.
+                self._rollbacks.append((m, fn, snapshot))
+        return n
+
     @property
     def batch_size(self):
         for c in self.cache:

mlx_lm/models/gpt_oss.py

@@ -230,6 +230,11 @@ def __call__(
 
 
 class Model(nn.Module):
+    # Sliding-window (RotatingKVCache) layers record an exact rollback while
+    # speculating, so the cache is trimmable within the verify window (see
+    # RotatingKVCache.record_rollback) — enables --draft-model speculation.
+    supports_speculative_rollback = True
+
     def __init__(self, args: ModelArgs):
         super().__init__()
         self.args = args