From 6eb374f56d636dcf9d00d5fb5a95ca7ba441def9 Mon Sep 17 00:00:00 2001 From: Pierre Lamy Date: Mon, 27 Jul 2026 12:38:23 -0400 Subject: [PATCH] fix: roll back MTP verify cache on early abort --- tests/test_mtp_spec_decode.py | 103 ++++++++++++++++++++++++++ vllm_mlx/spec_decode/mtp/generator.py | 49 +++++++----- 2 files changed, 134 insertions(+), 18 deletions(-) diff --git a/tests/test_mtp_spec_decode.py b/tests/test_mtp_spec_decode.py index 46df16227..77ac6ad99 100644 --- a/tests/test_mtp_spec_decode.py +++ b/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. diff --git a/vllm_mlx/spec_decode/mtp/generator.py b/vllm_mlx/spec_decode/mtp/generator.py index c56a9601b..5689a20ec 100644 --- a/vllm_mlx/spec_decode/mtp/generator.py +++ b/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,13 +795,25 @@ 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): @@ -798,7 +821,6 @@ def _record_round(k_used: int, round_wall_ms: float, accepts: list[bool]) -> Non # 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,7 +888,7 @@ 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 @@ -874,15 +896,6 @@ def _record_round(k_used: int, round_wall_ms: float, accepts: list[bool]) -> Non # _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