Skip to content

fix: roll back MTP verify cache on early abort - #1246

Draft
pierre427 wants to merge 1 commit into
raullenchai:mainfrom
pierre427:codex/speculative-rewind-guard
Draft

fix: roll back MTP verify cache on early abort#1246
pierre427 wants to merge 1 commit into
raullenchai:mainfrom
pierre427:codex/speculative-rewind-guard

Conversation

@pierre427

Copy link
Copy Markdown
Contributor

Summary

  • roll back the uncommitted target-cache draft suffix if an MTP verify round aborts during the single sync / host materialization boundary
  • centralize MTP verify-round rollback so the normal reject path trims the MTP draft cache through the same helper
  • add a regression test that injects a materialization fault after target verify advances caches but before fresh accept/reject state is read

Why

The target verify forward appends [current_token, draft_1, ..., draft_K] to the main model cache, while the MTP drafter has already appended the K proposed draft positions to its own cache. If cancellation, an injected fault, or a host materialization error fires before the fresh accept state is safely available, the generator should keep the already-committed current token and drop every uncommitted draft position from both caches.

This is intentionally narrow: normal all-accept behavior is unchanged, and the existing reject behavior now goes through the same combined rollback helper.

Validation

  • CPU-forced focused tests: tests/test_mtp_spec_decode.py::test_generator_rolls_back_verify_round_on_early_materialization_abort, ::test_generator_emits_first_token_from_backbone_then_draft, ::test_generator_rejection_path_does_not_count_as_accept — 3 passed
  • ruff check vllm_mlx/spec_decode/mtp/generator.py tests/test_mtp_spec_decode.py
  • ruff format --check vllm_mlx/spec_decode/mtp/generator.py tests/test_mtp_spec_decode.py
  • git diff --check -- vllm_mlx/spec_decode/mtp/generator.py tests/test_mtp_spec_decode.py

Full tests/test_mtp_spec_decode.py was attempted CPU-only in this worktree and is environment-blocked on the local system Python 3.9 / missing optional deps (fastapi, mlx_lm.models.qwen3_5), with the same broader toolchain mismatch outside this diff. Default-GPU validation was intentionally not run because a local Qwen3.5-122B quality gate is active on the machine.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant