[Core][Diffusion] Add pause_generation(mode="keep") for diffusion stages (#7685)

Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>

tests/diffusion/test_diffusion_worker_synchronize_device.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Unit tests for the device step of the diffusion pause barrier."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_worker(runner: object) -> DiffusionWorker:
+    worker = object.__new__(DiffusionWorker)
+    worker.model_runner = runner
+    worker.rank = 0
+    return worker
+
+
+def test_synchronize_device_waits_for_prefetch_before_device_sync(mocker):
+    platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    order: list = []
+    manager = MagicMock()
+    manager.wait_prefetch.side_effect = lambda timeout: order.append(("prefetch", timeout)) or True
+    platform.synchronize.side_effect = lambda: order.append("sync")
+    worker = _make_worker(SimpleNamespace(kv_transfer_manager=manager))
+
+    worker.synchronize_device(timeout=2.5)
+
+    assert order == [("prefetch", 2.5), "sync"]
+
+
+def test_synchronize_device_fails_when_prefetch_is_still_outstanding(mocker):
+    platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    manager = MagicMock()
+    manager.wait_prefetch.return_value = False
+    worker = _make_worker(SimpleNamespace(kv_transfer_manager=manager))
+
+    with pytest.raises(TimeoutError, match="prefetch"):
+        worker.synchronize_device(timeout=0.01)
+
+    platform.synchronize.assert_not_called()
+
+
+def test_synchronize_device_without_transfer_manager_only_syncs(mocker):
+    platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    worker = _make_worker(SimpleNamespace(kv_transfer_manager=None))
+
+    worker.synchronize_device(timeout=None)
+
+    platform.synchronize.assert_called_once_with()