fix: roll back MTP verify cache on early abort

tests/test_mtp_spec_decode.py

@@ -2192,6 +2192,55 @@ def make_mtp_cache(self):
         return []
 
 
+class _CountingKVCache:
+    """Tiny trimmable cache double for MTP rollback accounting tests."""
+
+    def __init__(self):
+        self.offset = 0
+        self.trim_calls: list[int] = []
+
+    def is_trimmable(self):
+        return True
+
+    def trim(self, n):
+        if n < 0:
+            raise AssertionError(f"negative trim: {n}")
+        self.trim_calls.append(n)
+        self.offset -= min(self.offset, n)
+
+
+class _CacheAdvancingQwen35Model(_MockedQwen35Model):
+    """Mock that advances supplied cache doubles on each forward."""
+
+    def __init__(self, backbone_outputs: list[int], mtp_outputs: list[int]):
+        super().__init__(backbone_outputs, mtp_outputs)
+        self.layers = [object()]
+
+    def __call__(
+        self,
+        inputs,
+        cache=None,
+        input_embeddings=None,
+        return_hidden: bool = False,
+        n_confirmed: int = 0,
+    ):
+        if cache is not None:
+            for c in cache:
+                c.offset += int(inputs.shape[1])
+        return super().__call__(
+            inputs,
+            cache=cache,
+            input_embeddings=input_embeddings,
+            return_hidden=return_hidden,
+            n_confirmed=n_confirmed,
+        )
+
+    def mtp_forward(self, hidden, next_token_ids, mtp_cache):
+        for c in mtp_cache:
+            c.offset += int(next_token_ids.shape[1])
+        return super().mtp_forward(hidden, next_token_ids, mtp_cache)
+
+
 def test_generator_emits_first_token_from_backbone_then_draft():
     """First yield comes from the backbone (``from_draft=False``); on
     accept the second yield is the MTP draft (``from_draft=True``).
@@ -2251,6 +2300,60 @@ def test_generator_emits_first_token_from_backbone_then_draft():
     assert snap.tokens_saved == 1
 
 
+def test_generator_rolls_back_verify_round_on_early_materialization_abort(
+    monkeypatch,
+):
+    """Abort after target verify advances caches, before accept state is fresh.
+
+    The first generator step commits the primary token and builds one MTP draft.
+    The second step runs the target verify forward over ``[primary, draft]``.
+    We then force the host sync/materialization boundary to raise. The guard
+    must keep the committed primary in the target cache while dropping the
+    uncommitted draft from both target and MTP caches.
+    """
+
+    import vllm_mlx.spec_decode.mtp.generator as generator_mod
+    from vllm_mlx.spec_decode.mtp.accept_counter import MTPAcceptCounter
+    from vllm_mlx.spec_decode.mtp.generator import mtp_generate_step
+
+    model_cache = _CountingKVCache()
+    mtp_cache = _CountingKVCache()
+    model = _CacheAdvancingQwen35Model([7, 11, 13], [11])
+    prompt = mx.array([1], dtype=mx.uint32)
+
+    gen = mtp_generate_step(
+        prompt,
+        model,
+        max_tokens=3,
+        prompt_cache=[model_cache, mtp_cache],
+        accept_counter=MTPAcceptCounter(),
+        disable_auto_k=True,
+    )
+
+    assert next(gen)[0] == 7
+    assert model_cache.offset == 1
+    # The draft is generated when the generator resumes for the next token.
+    assert mtp_cache.offset == 0
+
+    eval_calls = 0
+
+    def _boom_on_verify_sync(*_args, **_kwargs):
+        nonlocal eval_calls
+        eval_calls += 1
+        if eval_calls >= 2:
+            raise RuntimeError("sentinel materialization abort")
+
+    monkeypatch.setattr(generator_mod.mx, "eval", _boom_on_verify_sync)
+
+    with pytest.raises(RuntimeError, match="sentinel materialization abort"):
+        next(gen)
+
+    assert model_cache.offset == 2
+    assert model_cache.trim_calls[-1] == 1
+    assert mtp_cache.offset == 0
+    assert mtp_cache.trim_calls[-1] == 1
+
+
 def test_generator_rejection_path_does_not_count_as_accept():
     """When draft != verify_pred at temp=0 the generator takes the
     reject branch — counter shows attempt without accept.

vllm_mlx/spec_decode/mtp/generator.py

@@ -385,6 +385,17 @@ def _rollback_draft(n_to_drop: int = 1):
             elif c.is_trimmable():
                 c.trim(n_to_drop)
 
+    def _rollback_verify_round(n_to_drop: int) -> None:
+        """Roll back uncommitted target + MTP draft state for one verify round."""
+        if n_to_drop <= 0:
+            _clear_rollback()
+            return
+
+        _rollback_draft(n_to_drop)
+        for mc in mtp_cache:
+            if mc.is_trimmable():
+                mc.trim(n_to_drop)
+
     def _step_backbone(yy, prev, n_predict=1, n_confirmed=0, xtc_draw=None):
         """Run backbone on ``yy`` and return (tokens, logprobs, accept_lps, hidden, prev)."""
         with mx.stream(generation_stream):
@@ -784,21 +795,32 @@ def _record_round(k_used: int, round_wall_ms: float, accepts: list[bool]) -> Non
                 bonus_tok_arr = toks[k_len]
 
             # ------- SINGLE SYNC -------
-            mx.eval(toks, accept_mask_arr, residual_toks_arr, bonus_tok_arr, u)
-
-            # ------- Host-side read (all values already resident) -------
-            accept_flags = accept_mask_arr.tolist()
-            residual_ids = residual_toks_arr.tolist()
-            bonus_id = int(bonus_tok_arr.item())
-            draft_ids = drafts_arr.tolist()
+            accepted_count = 0
+            try:
+                mx.eval(toks, accept_mask_arr, residual_toks_arr, bonus_tok_arr, u)
+
+                # ------- Host-side read (all values already resident) -------
+                accept_flags = accept_mask_arr.tolist()
+                residual_ids = residual_toks_arr.tolist()
+                bonus_id = int(bonus_tok_arr.item())
+                draft_ids = drafts_arr.tolist()
+            except BaseException:
+                # The target verify forward above has already appended
+                # ``[y, d_1, ..., d_K]`` to model_cache, and the prior
+                # MTP chain appended the K draft positions to mtp_cache.
+                # If a cancellation / injected fault / host materialization
+                # error fires before this round's fresh accept count is known,
+                # never let stale state leak into a fallback path: keep the
+                # committed ``y`` position and drop every uncommitted draft.
+                _rollback_verify_round(k_len - accepted_count)
+                raise
 
             # Bump attempts by K (one per draft position considered).
             for _ in range(k_len):
                 accept_counter.record_attempt()
 
             # Sequential accept-reject walk (host-only; no MLX ops).
             accepts: list[bool] = []
-            accepted_count = 0
             for i in range(k_len):
                 ok = bool(accept_flags[i])
                 accepts.append(ok)
@@ -866,23 +888,14 @@ def _record_round(k_used: int, round_wall_ms: float, accepts: list[bool]) -> Non
                 # (k_len - accepted_count) unaccepted drafts from the
                 # caches.
                 n_to_drop = k_len - accepted_count
-                _rollback_draft(n_to_drop)
+                _rollback_verify_round(n_to_drop)
                 accept_counter.record_reject()
                 if logits_processors and prev_tokens is not None:
                     # Discard the ``n_to_drop`` rejected positions
                     # from prev_tokens (they were appended by
                     # _step_backbone during the batched verify).
                     prev_tokens = prev_tokens[:-n_to_drop]
 
-                # Also trim mtp_cache by the same n_to_drop — those
-                # positions were appended by _step_mtp_chain and
-                # correspond to the rejected drafts. The MTP KV
-                # cache is per-layer KVCache (see qwen3_5_inject
-                # make_mtp_cache / gemma4_inject) — always trimmable.
-                for mc in mtp_cache:
-                    if mc.is_trimmable():
-                        mc.trim(n_to_drop)
-
                 verify_tok_id = int(residual_ids[accepted_count])
 
                 ntoks += 1