[Bugfix] Stop cancelled MiniMax-H3 requests at denoising boundaries (#7844)

Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>

docs/serving/videos_api.md

@@ -120,7 +120,18 @@ request for execution.
 `DELETE /v1/videos/{video_id}` issues a bounded engine abort
 (`VLLM_OMNI_ABORT_TIMEOUT`, default 2s), then cancels the frontend
 task. Cancellation cleanup is also bounded and best-effort: it confirms
-the abort was queued, and the current request batch may still drain.
+the abort was submitted. In request execution mode, MiniMax-H3 checks
+cancellation at model boundaries,
+including before and after input preparation, after each denoising step, and
+before decode, and skips the remaining generation when cancelled. Input
+preparation or decoding already in progress may finish before the next boundary.
+Successful steps do not force device synchronization for cancellation.
+Parallel ranks agree before stopping; independent requests sharing a distributed
+AllGather offload wave can stop that wave early only when all its requests are
+cancelled, so a cancelled request cannot strand its live peers in a collective.
+In step execution mode, cancellation is handled by the scheduler between steps;
+the request-mode component-boundary checks do not apply within an active step.
+Other pipelines may still drain their current request batch.
 The job is then re-read so a completed save is not orphaned.
 
 ### Synchronous Response

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -3494,3 +3494,48 @@ def test_long_video_shape_requires_explicit_opt_in(duration):
     assert frames >= duration * 24 and frames % 17 == 5
     assert video_t == (frames - 5) // 17 * 5 + 2
     assert audio_t == round(frames / 24 * 40)
+
+
+@pytest.mark.parametrize("preencode", [False, True])
+@pytest.mark.parametrize("cancel_phase", ["before_prepare", "prepare", "diffuse"])
+def test_request_cancellation_at_prepare_and_decode_boundaries(preencode, cancel_phase, monkeypatch):
+    from vllm_omni.diffusion.cancellation import RequestCancellationRegistry, request_cancellation_scope
+    from vllm_omni.diffusion.data import DiffusionRequestAbortedError
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.platforms import current_omni_platform
+
+    synchronize = Mock()
+    monkeypatch.setattr(current_omni_platform, "synchronize", synchronize)
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    registry = RequestCancellationRegistry()
+    signal = registry.create("request")
+
+    def prepare(*args):
+        if cancel_phase == "prepare":
+            registry.cancel(["request"])
+        return {"num_outputs": 1, "seed": 1101, "preencode_mp4": preencode}
+
+    def diffuse(**kwargs):
+        registry.cancel(["request"])
+        return torch.zeros(1), torch.zeros(1)
+
+    pipeline._prepare_request_inputs = Mock(side_effect=prepare)
+    pipeline._denoise_kwargs = Mock(return_value={})
+    pipeline.diffuse = Mock(side_effect=diffuse)
+    pipeline.decode = Mock()
+    pipeline.decode_to_mp4 = Mock()
+    try:
+        if cancel_phase == "before_prepare":
+            registry.cancel(["request"])
+        with request_cancellation_scope([signal]), pytest.raises(DiffusionRequestAbortedError):
+            pipeline.forward(_t2va_batch())
+        synchronize.assert_called_once_with()
+        if cancel_phase == "before_prepare":
+            pipeline._prepare_request_inputs.assert_not_called()
+        if cancel_phase != "diffuse":
+            pipeline.diffuse.assert_not_called()
+        pipeline.decode.assert_not_called()
+        pipeline.decode_to_mp4.assert_not_called()
+    finally:
+        registry.close()

tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py

@@ -171,6 +171,41 @@ def test_step_execution_matches_request_mode_denoise_loop(num_steps, mocker):
     torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], reference_audio)
 
 
+def test_request_mode_cancellation_stops_before_next_denoise_step(monkeypatch):
+    from vllm_omni.diffusion.cancellation import RequestCancellationRegistry, request_cancellation_scope
+    from vllm_omni.diffusion.data import DiffusionRequestAbortedError
+    from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop
+    from vllm_omni.platforms import current_omni_platform
+
+    # This test runs real packing/Euler updates on CPU with the small DiT above.
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+    branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=5)
+    registry = RequestCancellationRegistry()
+    signal = registry.create("request")
+    steps = []
+
+    def cancel_after_first_step(step, video, audio):
+        steps.append(step)
+        registry.cancel(["request"])
+
+    try:
+        with request_cancellation_scope([signal]), pytest.raises(DiffusionRequestAbortedError):
+            minimax_h3_denoise_loop(
+                model=_SegmentMeanModel(),
+                positive=branch,
+                initial_video_rows=video_rows,
+                initial_audio_rows=audio_rows,
+                keyframe_cond_rows=None,
+                sigmas_video=_sigmas(6, 12.0),
+                sigmas_audio=_sigmas(6, 3.0),
+                device=torch.device("cpu"),
+                on_step=cancel_after_first_step,
+            )
+    finally:
+        registry.close()
+    assert steps == [0]
+
+
 def test_step_execution_matches_request_mode_with_latent_edits():
     """Masked model rows, row timesteps, and scheduler math match both paths."""
     from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod

tests/diffusion/test_diffusion_engine_cleanup.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import asyncio
 import queue
@@ -432,3 +432,68 @@ def test_finalize_aborted_request_without_async_output_skips_drop() -> None:
 
     assert output.aborted is True
     engine.executor.drop_output.assert_not_called()
+
+
+@pytest.mark.parametrize("scheduler_fails", [False, True])
+def test_fail_engine_releases_cancellation_signals_after_workers_stop(scheduler_fails):
+    from multiprocessing.shared_memory import SharedMemory
+
+    from vllm_omni.diffusion.cancellation import RequestCancellationRegistry
+
+    engine = _make_engine()
+    engine._fail_pending_rpcs = Mock()
+    registry = RequestCancellationRegistry()
+    engine._request_cancellations = registry
+    name = registry.create("running")
+
+    def shutdown():
+        # Cancellation is visible before shutdown, and late readers can attach.
+        reader = SharedMemory(name=name)
+        try:
+            assert reader.buf[0] == 1
+        finally:
+            reader.close()
+
+    engine.executor.shutdown.side_effect = shutdown
+    engine.scheduler.close = Mock(side_effect=RuntimeError("scheduler cleanup failed") if scheduler_fails else None)
+    try:
+        if scheduler_fails:
+            with pytest.raises(RuntimeError, match="scheduler cleanup failed"):
+                engine._fail_engine(RuntimeError("engine failed"))
+        else:
+            engine._fail_engine(RuntimeError("engine failed"))
+            engine.close()
+            engine.executor.shutdown.assert_called_once_with()
+        with pytest.raises(FileNotFoundError):
+            SharedMemory(name=name)
+        assert not registry._signals
+    finally:
+        registry.close()
+
+
+def test_close_retains_cancellation_signal_until_stuck_worker_stops():
+    from multiprocessing.shared_memory import SharedMemory
+
+    from vllm_omni.diffusion.cancellation import RequestCancellationRegistry
+
+    engine = _make_engine()
+    registry = RequestCancellationRegistry()
+    engine._request_cancellations = registry
+    name = registry.create("running")
+    engine.worker_thread = Mock()
+    engine.worker_thread.is_alive.side_effect = [True, True, False, False]
+    try:
+        engine.close()
+        engine.executor.shutdown.assert_not_called()
+        reader = SharedMemory(name=name)
+        try:
+            assert reader.buf[0] == 1
+        finally:
+            reader.close()
+        engine.close()
+        engine.executor.shutdown.assert_called_once_with()
+        with pytest.raises(FileNotFoundError):
+            SharedMemory(name=name)
+        assert not registry._signals
+    finally:
+        registry.close()

tests/diffusion/test_diffusion_model_runner.py

@@ -1411,3 +1411,44 @@ def test_vllm_set_forward_context_implementation(monkeypatch):
             ),
         ),
     ], ERROR_MESSAGE
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("cancel_all", [False, True])
+def test_execute_model_batch_cancellation_preserves_live_peer(monkeypatch, cancel_all):
+    from vllm_omni.diffusion.cancellation import (
+        RequestCancellationRegistry,
+        check_request_cancellation,
+    )
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+
+    class CancellableBatchPipeline(_BatchPipeline):
+        supports_request_cancellation = True
+
+        def forward(self, batch):
+            check_request_cancellation()
+            return super().forward(batch)
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
+    pipeline = CancellableBatchPipeline(outputs=[DiffusionOutput(output="a"), DiffusionOutput(output="b")])
+    runner = _make_batch_runner(pipeline)
+    sched = _make_scheduler_output(num_reqs=2)
+    registry = RequestCancellationRegistry()
+    try:
+        for entry in sched.scheduled_new_reqs:
+            entry.req.cancellation_signal = registry.create(entry.request_id)
+        registry.cancel(["req-0", "req-1"] if cancel_all else ["req-0"])
+        result = runner.execute_model_batch(sched, runner.od_config)
+        assert len(result.runner_outputs) == 2
+        assert [output.result.aborted for output in result.runner_outputs] == [cancel_all, cancel_all]
+        if not cancel_all:
+            assert result.runner_outputs[1].request_id == "req-1"
+            assert result.runner_outputs[1].result.output == "b"
+        else:
+            assert pipeline.last_batch is None
+    finally:
+        registry.close()

tests/diffusion/test_request_cancellation.py

@@ -0,0 +1,276 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Cancellation transport, ownership and collective agreement regressions."""
+
+import multiprocessing as mp
+import queue
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from multiprocessing.shared_memory import SharedMemory
+from unittest.mock import patch
+
+import pytest
+
+from vllm_omni.diffusion.cancellation import (
+    RequestCancellationRegistry,
+    check_request_cancellation,
+    request_cancellation_scope,
+)
+from vllm_omni.diffusion.data import DiffusionRequestAbortedError, OmniDiffusionConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.fixture
+def registry(monkeypatch):
+    from vllm_omni.platforms import current_omni_platform
+
+    # These tests exercise the cancellation protocol without device work.
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+    signals = RequestCancellationRegistry()
+    yield signals
+    signals.close()
+
+
+def test_cancel_is_request_scoped_and_id_reuse_is_fresh(registry):
+    first = registry.create("first")
+    other = registry.create("other")
+    registry.cancel(["missing", "first"])
+    with pytest.raises(ValueError, match="Duplicate cancellation request"):
+        registry.create("first")
+    with request_cancellation_scope([first]):
+        with pytest.raises(DiffusionRequestAbortedError):
+            check_request_cancellation()
+    with request_cancellation_scope([other]):
+        check_request_cancellation()
+    registry.finish("first")
+    registry.finish("first")
+    with pytest.raises(FileNotFoundError):
+        SharedMemory(name=first)
+    replacement = registry.create("first")
+    with request_cancellation_scope([replacement]):
+        check_request_cancellation()
+
+
+def test_scope_restores_outer_state_and_does_not_leak_to_threads(registry):
+    signal = registry.create("request")
+    registry.cancel_all()
+    with request_cancellation_scope([signal]):
+        with request_cancellation_scope([None]):
+            check_request_cancellation()
+        with ThreadPoolExecutor(max_workers=1) as executor:
+            executor.submit(check_request_cancellation).result(timeout=10)
+        with pytest.raises(DiffusionRequestAbortedError):
+            check_request_cancellation()
+    check_request_cancellation()
+
+
+def test_mixed_batch_keeps_uncancelled_request_running(registry):
+    first = registry.create("first")
+    other = registry.create("other")
+    with request_cancellation_scope([first, other]):
+        registry.cancel(["first"])
+        check_request_cancellation()
+        registry.cancel(["other"])
+        with pytest.raises(DiffusionRequestAbortedError):
+            check_request_cancellation()
+
+
+def test_checkpoint_observes_cancel_during_device_wait(registry, monkeypatch):
+    from vllm_omni.platforms import current_omni_platform
+
+    first = registry.create("first")
+    signal = registry.create("request")
+    registry.cancel(["first"])
+    # A local cancellation triggers the device wait. A peer can be cancelled
+    # during that wait; re-read before deciding whether the whole batch stops.
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: registry.cancel(["request"]))
+    with request_cancellation_scope([first, signal]):
+        with pytest.raises(DiffusionRequestAbortedError):
+            check_request_cancellation(synchronize=True)
+
+
+@pytest.mark.parametrize("synchronize_early", [False, True])
+def test_checkpoint_drains_once_before_abort_only(registry, monkeypatch, synchronize_early):
+    from unittest.mock import Mock
+
+    from vllm_omni.platforms import current_omni_platform
+
+    synchronize = Mock()
+    monkeypatch.setattr(current_omni_platform, "synchronize", synchronize)
+    signal = registry.create("request")
+    with request_cancellation_scope([signal]):
+        check_request_cancellation(synchronize=synchronize_early)
+        synchronize.assert_not_called()
+        # Simulate work queued after a clean checkpoint. Cancellation at the
+        # next checkpoint must drain it before exception unwinding drops refs.
+        pending_work = [object()]
+        synchronize.side_effect = pending_work.clear
+        registry.cancel(["request"])
+        with pytest.raises(DiffusionRequestAbortedError):
+            try:
+                check_request_cancellation(synchronize=synchronize_early)
+            finally:
+                assert not pending_work
+        synchronize.assert_called_once_with()
+
+
+def test_close_removes_all_owned_signals(registry):
+    names = [registry.create(str(i)) for i in range(3)]
+    registry.cancel_all()
+    registry.close()
+    registry.close()
+    for name in names:
+        with pytest.raises(FileNotFoundError):
+            SharedMemory(name=name)
+
+
+def _admission_engine(registry):
+    from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+    from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
+
+    engine = object.__new__(DiffusionEngine)
+    engine.scheduler = RequestScheduler()
+    engine.scheduler.initialize(OmniDiffusionConfig())
+    engine._request_cancellations = registry
+    engine._cv = threading.Condition(threading.RLock())
+    engine._closed = False
+    engine._out_streams = {}
+    engine.abort_queue = queue.Queue()
+    return engine
+
+
+def _request(request_id):
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    return OmniDiffusionRequest(
+        prompt="A fox in the snow.", sampling_params=OmniDiffusionSamplingParams(), request_id=request_id
+    )
+
+
+def test_engine_signals_worker_before_busy_loop_drains_abort(registry):
+    from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
+
+    engine = _admission_engine(registry)
+    request = _request("running")
+    engine._add_prepared_request(request)
+    engine.scheduler.schedule()
+    signal = request.cancellation_signal
+    assert signal is not None
+    with request_cancellation_scope([signal]):
+        engine.abort("running")
+        # The busy loop is still occupied by execution. Only the signal has
+        # changed; scheduler mutation must stay on that loop's thread.
+        assert engine.scheduler.get_request_state("running").status == DiffusionRequestStatus.RUNNING
+        with pytest.raises(DiffusionRequestAbortedError):
+            check_request_cancellation()
+    engine._process_aborts_queue()
+    assert engine._finalize_finished_request("running").aborted
+    assert engine.scheduler.get_request_state("running") is None
+    with pytest.raises(FileNotFoundError):
+        SharedMemory(name=signal)
+
+
+def test_rejected_admission_releases_signal(registry, monkeypatch):
+    engine = _admission_engine(registry)
+    allocated_names = []
+
+    def reject(request):
+        allocated_names.append(request.cancellation_signal)
+        raise ValueError("admission rejected")
+
+    monkeypatch.setattr(engine.scheduler, "add_request", reject)
+    request = _request("rejected")
+    with pytest.raises(ValueError, match="admission rejected"):
+        engine._add_prepared_request(request)
+    assert request.cancellation_signal is None
+    with pytest.raises(FileNotFoundError):
+        SharedMemory(name=allocated_names[0])
+    registry.create("rejected")
+
+
+def _collective_worker(rank, name, rendezvous, ready, first_check, second_check, results):
+    import torch
+    import torch.distributed as dist
+
+    from vllm_omni.diffusion.distributed import parallel_state
+    from vllm_omni.platforms import current_omni_platform
+
+    dist.init_process_group("gloo", init_method=f"file://{rendezvous}", rank=rank, world_size=2)
+    # Exercise a real CPU collective without constructing GPU model groups.
+    group = object.__new__(parallel_state.GroupCoordinator)
+    group.world_size = 2
+    group.cpu_group = dist.group.WORLD
+    parallel_state._WORLD = group
+    try:
+        with (
+            torch.inference_mode(),
+            patch.object(current_omni_platform, "synchronize") as synchronize,
+            request_cancellation_scope([name]),
+        ):
+            dist.barrier()
+            if rank == 0:
+                ready.set()
+            assert first_check.wait(120)
+            check_request_cancellation()
+            results.put((rank, "partial_wave_kept_running"))
+            # A replica missing its token (or all local requests) still votes
+            # false. Both ranks must reach each collective without deadlocking.
+            for absent in ([None], []):
+                with request_cancellation_scope([name] if rank == 0 else absent):
+                    check_request_cancellation()
+            results.put((rank, "missing_signal_kept_running"))
+            synchronize.assert_not_called()
+            assert second_check.wait(120)
+            try:
+                check_request_cancellation()
+            except DiffusionRequestAbortedError:
+                synchronize.assert_called_once_with()
+                results.put((rank, "fully_cancelled_wave_stopped"))
+            else:
+                raise AssertionError("All cancelled ranks must stop")
+    finally:
+        dist.destroy_process_group()
+
+
+def test_spawned_ranks_agree_before_leaving_collective_wave(registry, tmp_path):
+    ctx = mp.get_context("spawn")
+    names = [registry.create(str(rank)) for rank in range(2)]
+    ready, first_check, second_check = ctx.Event(), ctx.Event(), ctx.Event()
+    results = ctx.Queue()
+    workers = [
+        ctx.Process(
+            target=_collective_worker,
+            args=(rank, names[rank], str(tmp_path / "rendezvous"), ready, first_check, second_check, results),
+        )
+        for rank in range(2)
+    ]
+    for worker in workers:
+        worker.start()
+    try:
+        assert ready.wait(120)
+        registry.cancel(["0"])
+        first_check.set()
+        assert {results.get(timeout=120) for _ in range(2 * len(workers))} == {
+            (0, "partial_wave_kept_running"),
+            (1, "partial_wave_kept_running"),
+            (0, "missing_signal_kept_running"),
+            (1, "missing_signal_kept_running"),
+        }
+        registry.cancel(["1"])
+        second_check.set()
+        assert {results.get(timeout=120) for _ in workers} == {
+            (0, "fully_cancelled_wave_stopped"),
+            (1, "fully_cancelled_wave_stopped"),
+        }
+        for worker in workers:
+            worker.join(timeout=30)
+            assert worker.exitcode == 0
+    finally:
+        for worker in workers:
+            if worker.is_alive():
+                worker.terminate()
+            worker.join(timeout=30)
+        results.close()
+        results.join_thread()

vllm_omni/diffusion/cancellation.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Cooperative cancellation for local, full-forward diffusion workers.
+
+The engine owns each signal until the worker has returned. Unlike an executor
+RPC, writing the signal does not wait behind the forward being cancelled.
+Opted-in pipelines call ``check_request_cancellation`` at safe boundaries.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Iterable, Iterator, Sequence
+from contextlib import contextmanager
+from contextvars import ContextVar
+from multiprocessing.shared_memory import SharedMemory
+
+
+class RequestCancellationRegistry:
+    """Engine-owned signals; scheduler mutation remains on the engine thread."""
+
+    def __init__(self) -> None:
+        self._signals: dict[str, SharedMemory] = {}
+        self._lock = threading.Lock()
+
+    def create(self, request_id: str) -> str:
+        with self._lock:
+            if request_id in self._signals:
+                raise ValueError(f"Duplicate cancellation request: {request_id}")
+            signal = SharedMemory(create=True, size=1)
+            signal.buf[0] = 0
+            self._signals[request_id] = signal
+            return signal.name
+
+    def cancel(self, request_ids: Iterable[str]) -> None:
+        with self._lock:
+            for request_id in request_ids:
+                signal = self._signals.get(request_id)
+                if signal is not None:
+                    signal.buf[0] = 1
+
+    def cancel_all(self) -> None:
+        with self._lock:
+            for signal in self._signals.values():
+                signal.buf[0] = 1
+
+    @staticmethod
+    def _dispose(signal: SharedMemory) -> None:
+        try:
+            signal.unlink()
+        except FileNotFoundError:
+            # A process-level resource tracker may already have removed it.
+            pass
+        finally:
+            signal.close()
+
+    def finish(self, request_id: str) -> None:
+        """Release only after execution returns, including aborted execution."""
+        with self._lock:
+            signal = self._signals.pop(request_id, None)
+            if signal is not None:
+                self._dispose(signal)
+
+    def close(self) -> None:
+        """Release remaining signals after the executor has shut down."""
+        with self._lock:
+            for signal in self._signals.values():
+                self._dispose(signal)
+            self._signals.clear()
+
+
+_current_signals: ContextVar[tuple[SharedMemory | None, ...] | None] = ContextVar(
+    "diffusion_request_cancellation", default=None
+)
+
+
+@contextmanager
+def request_cancellation_scope(signal_names: Sequence[str | None], *, enabled: bool = True) -> Iterator[None]:
+    """Attach worker readers without taking ownership of the engine's names."""
+    signals: list[SharedMemory | None] = []
+    try:
+        for name in signal_names if enabled else ():
+            signals.append(SharedMemory(name=name) if name is not None else None)
+        # An opted-in rank without a local signal must still join every
+        # checkpoint collective and vote "not cancelled". None means only
+        # that this execution does not use cooperative cancellation at all.
+        token = _current_signals.set(tuple(signals) if enabled else None)
+        try:
+            yield
+        finally:
+            _current_signals.reset(token)
+    finally:
+        for signal in signals:
+            if signal is not None:
+                signal.close()
+
+
+def check_request_cancellation(*, synchronize: bool = False) -> None:
+    """Stop a cancelled execution wave without stranding a peer's collectives.
+
+    ``synchronize=True`` drains queued device work only when cancellation has
+    been requested locally, before re-reading the flags. Every abort drains queued
+    device work before unwinding. Successful steps retain asynchronous execution
+    when no local cancellation is observed.
+    Independent requests coupled by an AllGather offload wave must all be cancelled
+    before the wave can exit; cancelling one must not abort its live peers.
+    """
+    signals = _current_signals.get()
+    if signals is None:
+        return
+
+    import torch
+
+    locally_requested = any(signal is not None and signal.buf[0] for signal in signals)
+    device_drained = False
+    if synchronize and locally_requested:
+        from vllm_omni.platforms import current_omni_platform
+
+        current_omni_platform.synchronize()
+        device_drained = True
+
+    cancelled = bool(signals) and all(signal is not None and signal.buf[0] for signal in signals)
+    if torch.distributed.is_initialized():
+        from vllm_omni.diffusion.distributed.parallel_state import get_world_group
+
+        group = get_world_group()
+        if group.world_size > 1:
+            # Reading a shared byte alone is insufficient: ranks can observe a
+            # concurrent write on different steps. Make one collective decision.
+            flag = torch.tensor([int(cancelled)], dtype=torch.int32, device="cpu")
+            torch.distributed.all_reduce(flag, op=torch.distributed.ReduceOp.MIN, group=group.cpu_group)
+            cancelled = bool(flag.item())
+
+    if cancelled:
+        from vllm_omni.diffusion.data import DiffusionRequestAbortedError
+
+        if not device_drained:
+            from vllm_omni.platforms import current_omni_platform
+
+            current_omni_platform.synchronize()
+        raise DiffusionRequestAbortedError("Request cancelled at a model execution boundary")

vllm_omni/diffusion/diffusion_engine.py

@@ -25,6 +25,7 @@ from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.v1.engine.exceptions import EngineDeadError
 from vllm.v1.kv_cache_interface import KVCacheConfig
 
+from vllm_omni.diffusion.cancellation import RequestCancellationRegistry
 from vllm_omni.diffusion.data import (
     DIFFUSION_REQUEST_LIFECYCLE_KEY,
     DIFFUSION_REQUEST_STARTED,
@@ -180,6 +181,19 @@ def _max_num_seqs(od_config: OmniDiffusionConfig) -> int:
         return 1
 
 
+def supports_request_cancellation(od_config: OmniDiffusionConfig) -> bool:
+    """Whether the local pipeline checks cooperative cancellation boundaries."""
+    model_cls = _resolve_custom_pipeline_cls(getattr(od_config, "custom_pipeline_args", None))
+    if model_cls is None:
+        name = (
+            "DiffusersAdapterPipeline"
+            if uses_diffusers_adapter(od_config)
+            else getattr(od_config, "model_class_name", None)
+        )
+        model_cls = DiffusionModelRegistry._try_load_model_cls(name)
+    return getattr(model_cls, "supports_request_cancellation", False) is True
+
+
 def _uses_dlo_dp_concurrency(od_config: OmniDiffusionConfig) -> bool:
     parallel_config = getattr(od_config, "parallel_config", None)
     dp_size = getattr(parallel_config, "data_parallel_size", 1)
@@ -230,6 +244,8 @@ class DiffusionEngine:
     # don't hit AttributeError when _busy_loop accesses them.
     dp_concurrent: bool = False
     _scheduling_paused: bool = False
+    # Disabled until runtime initialization resolves the pipeline capability.
+    _request_cancellations: RequestCancellationRegistry | None = None
 
     def __init__(
         self,
@@ -393,6 +409,9 @@ class DiffusionEngine:
         self._shutting_down = False
         self._shutdown_complete = False
         self.abort_queue: queue.Queue[str] = queue.Queue()
+        self._request_cancellations = (
+            RequestCancellationRegistry() if supports_request_cancellation(self.od_config) else None
+        )
         self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue()
         # pause_scheduler(mode="keep"): no new batch is scheduled while set.
         self._scheduling_paused = False
@@ -913,6 +932,8 @@ class DiffusionEngine:
             # path, and repeating executor.shutdown() is unsafe.
             self._shutting_down = True
             self._closed = True
+            if self._request_cancellations is not None:
+                self._request_cancellations.cancel_all()
             if self.stop_event is not None:
                 self.stop_event.set()
             streams = list(self._out_streams.values())
@@ -925,7 +946,13 @@ class DiffusionEngine:
         # If Worker shutdown fails, retain the Scheduler reservations. A
         # remote producer may still be writing into those allocations.
         self.executor.shutdown()
-        self.scheduler.close()
+        try:
+            self.scheduler.close()
+        finally:
+            # Workers are down, so readers cannot attach or use the signals.
+            # Release them even if scheduler cleanup itself fails.
+            if self._request_cancellations is not None:
+                self._request_cancellations.close()
         self._shutdown_complete = True
 
     def _emit_finished_outputs(
@@ -1106,7 +1133,15 @@ class DiffusionEngine:
             if self._closed:
                 raise RuntimeError("DiffusionEngine is closed.")
             queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
-            request_id = self.scheduler.add_request(request)
+            if self._request_cancellations is not None:
+                request.cancellation_signal = self._request_cancellations.create(request.request_id)
+            try:
+                request_id = self.scheduler.add_request(request)
+            except BaseException:
+                if self._request_cancellations is not None:
+                    self._request_cancellations.finish(request.request_id)
+                    request.cancellation_signal = None
+                raise
             self._out_streams[request_id] = queue
             self._cv.notify_all()
 
@@ -1517,6 +1552,8 @@ class DiffusionEngine:
                 return
             if not self._closed:
                 self._closed = True
+                if self._request_cancellations is not None:
+                    self._request_cancellations.cancel_all()
                 if self.stop_event is not None:
                     self.stop_event.set()
                 pending_streams = list(self._out_streams.values())
@@ -1532,6 +1569,9 @@ class DiffusionEngine:
             if worker_thread.is_alive():
                 worker_thread.join(timeout=10)
             if worker_thread.is_alive():
+                # Keep cancellation names available: an in-flight worker may
+                # not have attached its readers yet. A later close releases
+                # them after execution and executor shutdown have completed.
                 logger.warning(
                     "Worker thread did not terminate within 10s; scheduler and executor shutdown will be deferred."
                 )
@@ -1550,6 +1590,8 @@ class DiffusionEngine:
         else:
             self.scheduler.close()
             self.executor.shutdown()
+        if self._request_cancellations is not None:
+            self._request_cancellations.close()
         self._shutdown_complete = True
 
     def abort(self, request_id: str | Iterable[str]) -> None:
@@ -1558,6 +1600,10 @@ class DiffusionEngine:
         with self._cv:
             if self._closed:
                 return
+            if self._request_cancellations is not None:
+                # Do not queue this behind the full-forward executor call.
+                # Scheduler state is still mutated only by the busy loop.
+                self._request_cancellations.cancel(request_ids)
             for req_id in request_ids:
                 self.abort_queue.put(req_id)
             self._cv.notify_all()
@@ -1601,6 +1647,8 @@ class DiffusionEngine:
         state = self.scheduler.get_request_state(request_id)
         popped_state = self.scheduler.pop_request_state(request_id)
         state = state or popped_state
+        if self._request_cancellations is not None:
+            self._request_cancellations.finish(request_id)
 
         if state is None:
             raise RuntimeError(f"Diffusion scheduler lost state for request {request_id}.")

vllm_omni/diffusion/models/minimax_h3/denoise_loop.py

@@ -16,6 +16,7 @@ from typing import Any
 import torch
 
 from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout, VideoTokenSpan
+from vllm_omni.diffusion.cancellation import check_request_cancellation
 from vllm_omni.diffusion.forward_context import (
     set_forward_context_denoise_step_idx,
     set_forward_context_denoise_timestep,
@@ -366,6 +367,7 @@ def minimax_h3_denoise_loop(
 
     num_steps = len(sigmas_video) - 1
     for step in range(num_steps):
+        check_request_cancellation()
         step_cm = step_profiler(step) if step_profiler is not None else nullcontext()
         with step_cm:
             s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1]
@@ -450,6 +452,7 @@ def minimax_h3_denoise_loop(
                 audio_rows[~audio_update] = audio_anchor  # per-step audio ref reset
             if on_step is not None:
                 on_step(step, video_rows, audio_rows)
+            check_request_cancellation(synchronize=True)
 
     minimax_h3_publish_denoise_progress(None, None, None)
     return video_rows, audio_rows

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -27,6 +27,7 @@ from vllm_omni.diffusion.cache.cachedit import (
     CacheDiTBackend,
     RequestScopedCacheDiTRuntime,
 )
+from vllm_omni.diffusion.cancellation import check_request_cancellation
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.parallel_state import get_world_group, init_world_group
 from vllm_omni.diffusion.distributed.utils import get_local_device
@@ -604,6 +605,7 @@ class MiniMaxH3Pipeline(
     """CFG-distilled joint video/audio generation for MiniMax H3."""
 
     supports_step_execution: ClassVar[bool] = True
+    supports_request_cancellation: ClassVar[bool] = True
 
     _dit_modules: ClassVar[list[str]] = ["transformer", "transformers_ref"]
     _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
@@ -2619,15 +2621,18 @@ class MiniMaxH3Pipeline(
     def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput:
         if len(request.prompts) != 1:
             raise OmniClientError("MiniMax H3 supports one request at a time")
+        check_request_cancellation()
         context = self._prepare_request_inputs(
             request.prompts[0],
             request.sampling_params,
         )
+        check_request_cancellation()
         denoise_kwargs = self._denoise_kwargs(context)
         num_outputs = context["num_outputs"]
         videos = []
         audios = []
         for output_seed in _minimax_h3_output_seeds(context["seed"], num_outputs):
+            check_request_cancellation()
             output_kwargs = {**denoise_kwargs, "seed": output_seed}
             if context.get("continuation") is None:
                 video_latent, audio_latent = self.diffuse(**output_kwargs)
@@ -2640,6 +2645,7 @@ class MiniMaxH3Pipeline(
                     overlap_frames=overlap_frames,
                     text_conditioning=context.get("continuation_text_conditioning"),
                 )
+            check_request_cancellation()
             if context["preencode_mp4"]:
                 videos.append(
                     self.decode_to_mp4(

vllm_omni/diffusion/request.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 import random
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
@@ -81,6 +81,9 @@ class OmniDiffusionRequest:
     kv_recv_ms: float = 0.0
     # Time spent waiting for initial admission by the diffusion scheduler.
     scheduler_queue_wait_ms: float | None = None
+    # Engine-owned shared-memory signal for cooperative full-forward workers.
+    # Only the engine creates/unlinks it; workers attach readers until return.
+    cancellation_signal: str | None = field(default=None, repr=False)
 
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -31,8 +31,9 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import (
     resolve_prompt_embed_cache_config,
 )
 from vllm_omni.diffusion.cache.selector import get_cache_backend
+from vllm_omni.diffusion.cancellation import check_request_cancellation, request_cancellation_scope
 from vllm_omni.diffusion.compile import regionally_compile
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError, OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
 from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
@@ -853,15 +854,29 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
                     in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
                 ),
                 paged_kv_context,
+                request_cancellation_scope(
+                    [getattr(req, "cancellation_signal", None) for req in reqs],
+                    enabled=getattr(self.pipeline, "supports_request_cancellation", False) is True,
+                ),
             ):
                 with record_function(record_name):
-                    raw_outputs = self.pipeline.forward(batch)
-                    outputs = _normalize_pipeline_outputs(
-                        raw_outputs,
-                        expected_count=len(reqs),
-                        allow_single_output=allow_single_output,
-                        pipeline_name=type(self.pipeline).__name__,
-                    )
+                    try:
+                        check_request_cancellation()
+                        raw_outputs = self.pipeline.forward(batch)
+                        outputs = _normalize_pipeline_outputs(
+                            raw_outputs,
+                            expected_count=len(reqs),
+                            allow_single_output=allow_single_output,
+                            pipeline_name=type(self.pipeline).__name__,
+                        )
+                    except DiffusionRequestAbortedError as exc:
+                        # The checkpoint aborts only a fully cancelled wave;
+                        # a mixed batch must keep running for its live peers.
+                        logger.info(
+                            "Stopped cancelled diffusion request(s) %s at a model execution boundary",
+                            [req.request_id for req in reqs],
+                        )
+                        outputs = [DiffusionOutput(aborted=True, abort_message=str(exc)) for _ in reqs]
                 with record_function("prepare_output_for_transport"):
                     outputs = [
                         self._prepare_output_for_transport(output, req.sampling_params)