[Feat][sleepmode] add omni sleepmode and ack protocol (#2022)

.buildkite/test-amd.yaml

@@ -117,3 +117,17 @@ steps:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+
+
+- label: "Omni Sleep Mode Test"
+  timeout_in_minutes: 40
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+    - pytest -s -v tests/e2e/offline_inference/test_omni_sleep_mode.py -m "advanced_model and omni and MI325" --run-level "advanced_model"

.buildkite/test-merge.yml

@@ -360,6 +360,46 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
+  - label: "Omni Sleep Mode Test with H100"
+    timeout_in_minutes: 30
+    depends_on: upload-merge-pipeline
+    commands:
+      - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+      - pytest -s -v tests/e2e/offline_inference/test_omni_sleep_mode.py -m "advanced_model and H100 and omni" --run-level "advanced_model"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
   - label: "Voxtral-TTS E2E Test"
     timeout_in_minutes: 20
     depends_on: upload-merge-pipeline

docs/features/sleep_mode.md

@@ -1,39 +1,243 @@
-# Sleep Mode
+# Sleep Mode & ACK Protocol
 
-vLLM-Omni’s **Sleep Mode** allows you to temporarily release most GPU memory used by a model—such as model weights and key-value (KV) caches (for autoregressive models)—**without stopping the server or unloading the Docker container**.
+vLLM-Omni’s **Sleep Mode** allows you to temporarily release most GPU memory used by a model—such as model weights and key-value (KV) caches—**without stopping the server or unloading the Docker container**.
 
-This feature is inherited from [vLLM’s Sleep Mode](https://blog.vllm.ai/2025/10/26/sleep-mode.html), which provides zero-reload model switching for multi-model serving.  
-
-It is especially useful in **RLHF**, **training**, or **cost-saving scenarios**, where GPU resources must be freed between inference workloads.
+This feature is inherited from [vLLM’s Sleep Mode](https://blog.vllm.ai/2025/10/26/sleep-mode.html) and extended with the **Omni ACK Protocol** to support multi-stage pipelines and heterogeneous hardware backends (NVIDIA, AMD, Intel, Huawei). It is especially useful in **RLHF**, **dynamic model switching**, or **cost-saving scenarios**.
 
 ---
 
-## Omni Model
+## 1. Feature Documentation
 
-Omni model inherit the feature from vLLM' Sleep Mode
+### Overview
+Omni Sleep Mode provides a mechanism to "sleep" specific model stages. When a stage enters sleep, its physical VRAM is reclaimed by the system, while the process state is preserved for rapid "wake-up" without full re-initialization.
 
-This means:
+### Sleep Levels
+We support two levels of hibernation to balance recovery speed and memory efficiency:
 
-- Support both Level 1 and Level 2 sleep, allow to release and reset both model weights and KV Cache
+| Level | Name | Mechanism | Recovery Speed | Memory Freed |
+| :--- | :--- | :--- | :--- | :--- |
+| **Level 1** | **Weight Offloading** | Offloads weights to Host CPU RAM. | **Fast** (DMA) | Substantial |
+| **Level 2** | **Full De-mapping** | Physically releases memory pages via VRAM scavenging. | **Moderate** | **Maximum** (up to 95%+) |
 
-## Diffusion Model Extension
+### Supported Platforms
 
-We added Sleep Mode support for **diffusion models**, which previously lacked this functionality.  
-In diffusion pipelines, this currently only offloads **model weight memory**, as these models typically do not use KV caches.
+Omni Sleep Mode is optimized for high-performance computing backends:
 
-This means:
+* **NVIDIA**: Supported via Virtual Memory Management (VMM).
+* **AMD (ROCm)**: Fully supported with physical page de-mapping.
+* **Intel XPU**: Supported via Level Zero memory management.
+* **Huawei NPU**: Supported via Ascend memory scavenging.
 
-- Diffusion models can now enter Level 1 sleep.
-- Pipeline states (e.g., noise schedulers, buffers) remain intact after waking.
-- Useful for releasing VRAM between image generation or training cycles.
+### Hardware Requirements
+* **Memory Considerations**: System RAM must be sufficient to hold offloaded weights during sleep.
+* **TP Support**: Tensor Parallel groups synchronize sleep/wake transitions across all workers.
 
 ---
 
-## Enable sleep mode
-To enable sleep mode, set the `enable_sleep_mode` in `engine_args` to `True`
 
+## 2. Usage Examples
+
+### Python API Example
+You can programmatically control the lifecycle of stages using the `AsyncOmni` engine.
 
-Example:
 ```python
-omni = Omni(model=...,enable_sleep_mode=True)
+
+import asyncio
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+async def run_sleep_demo():
+    # 1. initialization
+    engine = AsyncOmni(
+        model="ByteDance-Seed/BAGEL-7B-MoT",
+        enable_sleep_mode=True
+    )
+
+    # 2. sleep mode level2
+    acks = await engine.sleep(stage_ids=[0], level=2)
+    print(f"Freed {acks[0].freed_bytes / 1024**3:.2f} GiB on Stage 0")
+
+    # 3. wake up
+    await engine.wake_up(stage_ids=[0])
+
+if __name__ == "__main__":
+    asyncio.run(run_sleep_demo())
+
 ```
+
+### server command Example
+Start the server with sleep mode enabled:
+
+The first method
+
+```
+
+vllm serve ByteDance-Seed/BAGEL-7B-MoT \
+--omni \
+--enable-sleep-mode \
+--trust-remote-code \
+--gpu-memory-utilization 0.7
+
+```
+
+The second method
+
+```
+
+python3 -m vllm_omni.entrypoints.openai.api_server \
+    --model ByteDance-Seed/BAGEL-7B-MoT \
+    --omni \
+    --enable-sleep-mode \
+    --trust-remote-code \
+--gpu-memory-utilization 0.7
+
+```
+
+
+
+
+### Test Scenarios & Commands
+
+#### Scenario 1: LLM Engine Sleep
+
+Objective: Verify VRAM reclamation for Stage 0 (Thinker).
+
+Trigger sleep (Level 1 or Level 2) via client:
+
+```
+
+curl -X POST http://localhost:8000/v1/omni/sleep \
+     -H "Content-Type: application/json" \
+     -d '{"stage_ids": [0], "level": 2}'
+
+```
+
+Tip: Open a new terminal and run rocm-smi or nvidia-smi or to observe the immediate drop in VRAM usage.
+
+
+
+#### Scenario 2: Diffusion Sleep
+Objective: Verify VRAM reclamation for Stage 1 (Diffusion).
+
+Trigger sleep (Level 1 or Level 2) via client:
+
+```
+
+curl -X POST http://localhost:8000/v1/omni/sleep \
+     -H "Content-Type: application/json" \
+     -d '{"stage_ids": [1], "level": 2}'
+
+```
+
+
+
+#### Scenario 3: Multi-Stage Coordinated Stress Test
+Objective: Test concurrent sleep and rapid wake-up across multiple stages.
+
+Concurrent Sleep (Stage 0 & 1):
+
+```
+
+curl -X POST http://localhost:8000/v1/omni/sleep \
+     -H "Content-Type: application/json" \
+     -d '{"stage_ids": [0, 1], "level": 2}'
+
+```
+
+
+Rapid Wake-up:
+
+```
+
+curl -X POST http://localhost:8000/v1/omni/wakeup \
+     -H "Content-Type: application/json" \
+     -d '{"stage_ids": [0, 1]}'
+
+```
+
+
+#### Scenario 4: Full Lifecycle Memory Audit & Functional Integrity
+Objective: Audit the complete flow from Sleep to Wake-up followed by an Inference validation.
+
+Check Initial State: Observe baseline VRAM usage.
+
+Trigger Deep Sleep (Level 2):
+
+```
+
+curl -X POST http://localhost:8000/v1/omni/sleep \
+     -H "Content-Type: application/json" \
+     -d '{"stage_ids": [0], "level": 2}'
+
+```
+
+Wake-up Model:
+
+```
+
+curl -X POST http://localhost:8000/v1/omni/wakeup \
+     -H "Content-Type: application/json" \
+     -d '{"stage_ids": [0]}'
+
+```
+
+Verify Functional Integrity (Inference):
+Ensure the model still generates valid output after reloading weights.
+
+```
+
+curl -X POST http://localhost:8000/v1/images/generations \
+     -H "Content-Type: application/json" \
+     -d '{
+       "prompt": "A huge swimming pool, with many people swimming.",
+       "model": "ByteDance-Seed/BAGEL-7B-MoT",
+       "response_format": "b64_json",
+       "extra_body": {"sampling_params": {"num_inference_steps": 4, "seed": 42}}
+     }' > post.json
+
+```
+
+
+
+
+## 3. API Reference
+
+
+### Methods
+
+| Method | Arguments | Return Type | Description |
+| :--- | :--- | :--- | :--- |
+| **sleep** | `stage_ids: List[int], level: int` | `List[OmniACK]` | Triggers hibernation for specified stages. |
+| **wake_up** | `stage_ids: List[int]` | `List[OmniACK]` | Reloads weights and re-maps memory. |
+
+
+
+### OmniACK Dataclass Fields
+
+| Field | Type | Description |
+| :--- | :--- | :--- |
+| **task_id** | `str` | Unique identifier for the operation. |
+| **status** | `str` | `SUCCESS` or `ERROR`. |
+| **stage_id** | `int` | The ID of the stage that responded. |
+| **rank** | `int` | The rank ID within the Tensor Parallel group. |
+| **freed_bytes** | `int` | Actual amount of physical VRAM reclaimed. |
+| **metadata** | `dict` | Additional platform-specific metrics. |
+
+Metadata Field Analysis
+The metadata field is a dynamic dictionary containing hardware-specific telemetry and audit data, primarily used for verifying memory reclamation on various backends (e.g., AMD ROCm, NVIDIA CUDA).
+
+```
+"metadata": {
+    "source": "Platform_AMD_Instinct_MI300X",
+    "total_freed_gib": "78.57",
+    "rank_residual_gib": "2.07"
+}
+```
+
+#### Core Utility:
+**VRAM Reclamation Audit (total_freed_gib)**: Converts raw freed_bytes into human-readable GiB. It serves as the primary metric to verify that Level 2 sleep has successfully purged model weights from VRAM.
+
+**Residual & Fragmentation Monitoring (rank_residual_gib)**: Reports the remaining VRAM footprint after memory de-mapping. A low residual value (e.g., 2.07 GiB) confirms a successful "clean" state, ensuring the device is ready for high-memory co-located tasks like training or diffusion pipelines.
+
+**Backend Traceability (source)**: Identifies the underlying hardware driver or audit source. This is critical for debugging synchronization issues in multi-stage, distributed environments.
+
+**Performance Analytics (Roadmap)**: Future updates will include latency_ms (context-switch overhead) and cuda_graph_recalled (graph engine status) to optimize performance in high-frequency sleep/wake scenarios.

tests/diffusion/test_diffusion_worker.py

@@ -16,7 +16,7 @@ from pytest_mock import MockerFixture
 
 from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
 
 
 @pytest.fixture
@@ -81,17 +81,31 @@ class TestDiffusionWorkerLoadWeights:
 class TestDiffusionWorkerSleep:
     """Test DiffusionWorker.sleep method."""
 
+    @pytest.fixture(autouse=True)
+    def setup_allocator(self, mocker: MockerFixture):
+        """
+        Unified interception of Allocators, and provision of default security values.
+        """
+        self.mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+        self.mock_allocator = mocker.Mock()
+        self.mock_allocator_class.get_instance.return_value = self.mock_allocator
+        self.mock_allocator.get_current_usage.return_value = 4 * 1024**3
+        self.mock_allocator.sleep = mocker.Mock()
+
     def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):
         """Test sleep mode level 1 (offload weights only)."""
         mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
-        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3]
+        mock_platform.get_device_total_memory.return_value = 80 * 1024**3
         mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
 
         # Setup process-scoped memory mocks
         # Before sleep: 3GB used
         # After sleep: 1GB used (freed 2GB)
+        initial_usage = 3 * 1024**3
         mock_get_process_memory.side_effect = [
-            3 * 1024**3,
+            initial_usage,
             1 * 1024**3,
         ]
 
@@ -99,25 +113,29 @@ class TestDiffusionWorkerSleep:
         mock_allocator = mocker.Mock()
         mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
         mock_allocator.sleep = mocker.Mock()
+        mock_allocator.get_current_usage.return_value = initial_usage
 
         # Call sleep with level 1
         result = mock_gpu_worker.sleep(level=1)
 
         # Verify sleep was called with correct tags
         mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",))
-        assert result is True
+        assert bool(result) is True
         # Verify buffers were NOT saved (level 1 doesn't save buffers)
         assert len(mock_gpu_worker._sleep_saved_buffers) == 0
 
     def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker):
         """Test sleep mode level 2 (offload all, save buffers)."""
         mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
-        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_platform.get_free_memory.side_effect = [5 * 1024**3, 10 * 1024**3]
+        mock_platform.get_device_total_memory.return_value = 80 * 1024**3
         mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
 
         # Setup process-scoped memory mocks
+        initial_usage = 5 * 1024**3
         mock_get_process_memory.side_effect = [
-            5 * 1024**3,  # Before sleep
+            initial_usage,  # Before sleep
             1 * 1024**3,  # After sleep (freed 4GB)
         ]
 
@@ -125,6 +143,7 @@ class TestDiffusionWorkerSleep:
         mock_allocator = mocker.Mock()
         mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
         mock_allocator.sleep = mocker.Mock()
+        mock_allocator.get_current_usage.return_value = initial_usage
 
         # Mock pipeline buffers
         mock_buffer1 = torch.randn(10, 10)
@@ -141,7 +160,7 @@ class TestDiffusionWorkerSleep:
 
         # Verify sleep was called with empty tags (offload all)
         mock_allocator.sleep.assert_called_once_with(offload_tags=tuple())
-        assert result is True
+        assert bool(result) is True
 
         # Verify buffers were saved
         assert len(mock_gpu_worker._sleep_saved_buffers) == 2
@@ -151,22 +170,26 @@ class TestDiffusionWorkerSleep:
     def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker):
         """Test that sleep validates memory was actually freed."""
         mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
-        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_platform.get_free_memory.return_value = 10 * 1024**3
+        mock_platform.get_device_total_memory.return_value = 80 * 1024**3
         mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
 
         # Simulate process memory increase (should trigger assertion error)
+        initial_usage = 1 * 1024**3
         mock_get_process_memory.side_effect = [
-            1 * 1024**3,  # Before sleep: 1GB used
+            initial_usage,  # Before sleep: 1GB used
             3 * 1024**3,  # After sleep: 3GB used (negative freed)
         ]
 
         mock_allocator = mocker.Mock()
         mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
         mock_allocator.sleep = mocker.Mock()
+        mock_allocator.get_current_usage.return_value = initial_usage
 
         # This should raise an assertion error
-        with pytest.raises(AssertionError, match="Memory usage increased after sleeping"):
-            mock_gpu_worker.sleep(level=1)
+        result = mock_gpu_worker.sleep(level=1)
+        assert result == initial_usage
 
     def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(self, mocker: MockerFixture, mock_gpu_worker):
         """Test sleep uses device-scoped fallback when NVML is unavailable."""
@@ -184,11 +207,12 @@ class TestDiffusionWorkerSleep:
         mock_allocator = mocker.Mock()
         mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
         mock_allocator.sleep = mocker.Mock()
+        mock_allocator.get_current_usage.return_value = 2 * 1024**3
 
         result = mock_gpu_worker.sleep(level=1)
 
         mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",))
-        assert result is True
+        assert bool(result) is True
 
 
 class TestDiffusionWorkerWakeUp:
@@ -202,6 +226,7 @@ class TestDiffusionWorkerWakeUp:
         mock_allocator = mocker.Mock()
         mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
         mock_allocator.wake_up = mocker.Mock()
+        mock_allocator.get_current_usage.return_value = 10 * 1024**3
 
         # Ensure no saved buffers
         mock_gpu_worker._sleep_saved_buffers = {}
@@ -211,7 +236,7 @@ class TestDiffusionWorkerWakeUp:
 
         # Verify allocator.wake_up was called
         mock_allocator.wake_up.assert_called_once_with(["weights"])
-        assert result is True
+        assert bool(result) is True
 
     def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker):
         """Test wake_up with saved buffers (level 2 sleep)."""
@@ -221,6 +246,7 @@ class TestDiffusionWorkerWakeUp:
         mock_allocator = mocker.Mock()
         mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
         mock_allocator.wake_up = mocker.Mock()
+        mock_allocator.get_current_usage.return_value = 10 * 1024**3
 
         # Create saved buffers
         saved_buffer1 = torch.randn(10, 10)
@@ -255,7 +281,7 @@ class TestDiffusionWorkerWakeUp:
 
         # Verify saved buffers were cleared
         assert len(mock_gpu_worker._sleep_saved_buffers) == 0
-        assert result is True
+        assert bool(result) is True
 
     def test_wake_up_partial_buffer_restore(self, mocker: MockerFixture, mock_gpu_worker):
         """Test wake_up only restores buffers that were saved."""
@@ -265,6 +291,7 @@ class TestDiffusionWorkerWakeUp:
         mock_allocator = mocker.Mock()
         mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
         mock_allocator.wake_up = mocker.Mock()
+        mock_allocator.get_current_usage.return_value = 10 * 1024**3
 
         # Only save buffer1, not buffer2
         saved_buffer1 = torch.randn(10, 10)
@@ -293,4 +320,4 @@ class TestDiffusionWorkerWakeUp:
         # buffer2 should NOT be restored since it wasn't saved
         mock_buffer2.data.copy_.assert_not_called()
 
-        assert result is True
+        assert bool(result) is True

tests/e2e/offline_inference/test_omni_sleep_mode.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end tests for Omni Sleep Mode across various model architectures.
+"""
+
+import pytest
+import torch
+from vllm import SamplingParams
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
+MODEL_DIFF = "riverclouds/qwen_image_random"
+
+
+def get_ack_info(ack, key, default=None):
+    if hasattr(ack, key):
+        return getattr(ack, key)
+    if isinstance(ack, dict):
+        return ack.get(key, default)
+    return default
+
+
+def get_dynamic_devices(stage_idx, num_stages, tp_size):
+    total_gpus = torch.cuda.device_count()
+    gpus_per_stage = tp_size
+    start_idx = stage_idx * gpus_per_stage
+    if start_idx + gpus_per_stage > total_gpus:
+        start_idx = start_idx % total_gpus
+    device_ids = [str(start_idx + i) for i in range(gpus_per_stage)]
+    return ",".join(device_ids)
+
+
+# Test 1: Diffusion Model (2-Stage BAGEL)
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@pytest.mark.parametrize("tp_size", [1, 2])
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.asyncio
+async def test_diffusion_model_sleep_tp(tp_size):
+    num_gpus = torch.cuda.device_count()
+    if num_gpus < tp_size:
+        pytest.skip(f"Skipping TP={tp_size}")
+
+    engine_args = {
+        "model": MODEL,
+        "enable_sleep_mode": True,
+        "tensor_parallel_size": tp_size,
+        "enforce_eager": True,
+        "trust_remote_code": True,
+        "dtype": "bfloat16",
+        "gpu_memory_utilization": 0.5,
+    }
+
+    engine = AsyncOmni(**engine_args, stage_init_timeout=1200)
+    try:
+        # BAGEL requires 2 params
+        diff_sp = OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256)
+        llm_sp = SamplingParams()
+
+        # Warmup
+        async for _ in engine.generate("test", sampling_params=[llm_sp, diff_sp]):
+            pass
+
+        # Sleep all
+        acks = await engine.sleep(level=2)
+        statuses = [get_ack_info(ack, "status") for ack in acks]
+        assert all(s == "SUCCESS" for s in statuses), f"Sleep failed. Statuses: {statuses}"
+
+        # Wakeup & Verify
+        await engine.wake_up()
+        async for _ in engine.generate("verify", sampling_params=[llm_sp, diff_sp]):
+            pass
+
+        print(f"Diffusion TP={tp_size} Lifecycle OK")
+    finally:
+        engine.shutdown()
+
+
+# Test 2: Multi-stage Manual Config
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@pytest.mark.parametrize("tp_size", [1, 2])
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.asyncio
+async def test_multistage_sleep_h100(tp_size):
+    num_gpus = torch.cuda.device_count()
+    if num_gpus < tp_size * 2:
+        pytest.skip("Not enough GPUs")
+
+    stages = []
+    for i in range(2):
+        devs = get_dynamic_devices(i, 2, tp_size)
+        stages.append(
+            {
+                "stage_id": i,
+                "stage_type": "llm" if i == 0 else "diffusion",
+                "runtime": {"process": True, "devices": devs},
+                "engine_args": {
+                    "model": MODEL,
+                    "model_stage": "thinker" if i == 0 else "base",
+                    "tensor_parallel_size": tp_size,
+                    "gpu_memory_utilization": 0.4,
+                    "dtype": "bfloat16",
+                    "enable_sleep_mode": True,
+                    "trust_remote_code": True,
+                },
+            }
+        )
+
+    connectors = [{"src_stage_id": 0, "dst_stage_id": 1, "connector_type": "queue"}]
+
+    engine = AsyncOmni(
+        model=MODEL, stages=stages, connectors=connectors, enable_sleep_mode=True, stage_init_timeout=1200
+    )
+    try:
+        sp = OmniDiffusionSamplingParams(num_inference_steps=2)
+        async for _ in engine.generate("warmup", sampling_params=[SamplingParams(), sp]):
+            pass
+
+        acks = await engine.sleep(stage_ids=[0, 1], level=2)
+        assert len(acks) == 2 * tp_size
+
+        await engine.wake_up(stage_ids=[0, 1])
+        async for _ in engine.generate("verify", sampling_params=[SamplingParams(), sp]):
+            pass
+    finally:
+        engine.shutdown()
+
+
+# Test 3: Pure Diffusion Single-Stage
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@pytest.mark.parametrize("tp_size", [1, 2])
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.asyncio
+async def test_pure_diffusion_scenario(tp_size):
+    engine_args = {
+        "model": MODEL_DIFF,
+        "enable_sleep_mode": True,
+        "tensor_parallel_size": tp_size,
+        "enforce_eager": True,
+        "dtype": "bfloat16",
+        "gpu_memory_utilization": 0.5,
+    }
+
+    engine = AsyncOmni(**engine_args, stage_init_timeout=1200)
+    try:
+        await engine.sleep(level=1)
+        await engine.wake_up()
+        async for _ in engine.generate("test", sampling_params=[SamplingParams()]):
+            pass
+        print("Pure Diffusion OK")
+    finally:
+        engine.shutdown()

tests/entrypoints/test_omni_sleep_mode.py

@@ -0,0 +1,336 @@
+import asyncio
+import logging
+import os
+
+import pytest
+import torch
+from vllm import SamplingParams
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("OmniTest")
+pytestmark = [pytest.mark.advanced_model]
+
+
+def clean_gpu_envs():
+    """clean up GPU environment variables to ensure tests run on all available devices."""
+    device_visibility_vars = [
+        "CUDA_VISIBLE_DEVICES",  # NVIDIA
+        "HIP_VISIBLE_DEVICES",  # AMD ROCm
+        "ZE_AFFINITY_MASK",  # Intel XPU
+        "ONEAPI_DEVICE_SELECTOR",  # Intel OneAPI
+        "ASCEND_RT_VISIBLE_DEVICES",  # Huawei NPU (CAN)
+    ]
+    for key in device_visibility_vars:
+        os.environ.pop(key, None)
+
+
+def get_vram_info(device_id: int) -> dict:
+    """Obtain a snapshot of the specified GPU's memory (GiB)."""
+    try:
+        if current_omni_platform.is_rocm():
+            num_gpus = torch.cuda.device_count()
+            safe_id = device_id if device_id < num_gpus else 0
+            torch.cuda.synchronize(safe_id)
+            return {
+                "reserved": torch.cuda.memory_reserved(safe_id) / 1024**3,
+                "allocated": torch.cuda.memory_allocated(safe_id) / 1024**3,
+            }
+        else:
+            with torch.cuda.device(device_id):
+                torch.cuda.synchronize()
+                return {
+                    "reserved": torch.cuda.memory_reserved() / 1024**3,
+                    "allocated": torch.cuda.memory_allocated() / 1024**3,
+                }
+    except Exception as e:
+        logger.warning(f"memory skip ({device_id}): {e}")
+        return {"reserved": 0.0, "allocated": 0.0}
+
+
+def get_ack_info(ack, key, default=None):
+    """
+    Since ACKs in a distributed environment can be either objects or dictionaries,
+    this tool ensures compatibility.
+    """
+    if hasattr(ack, key):
+        return getattr(ack, key)
+    if isinstance(ack, dict):
+        return ack.get(key, default)
+    return default
+
+
+@pytest.fixture(scope="function")
+async def llm_engine():
+    if current_omni_platform.is_rocm():
+        clean_gpu_envs()
+    model_name = "ByteDance-Seed/BAGEL-7B-MoT"
+    common_args = {
+        "worker_type": "ar",
+        "enable_sleep_mode": True,
+        "dtype": "bfloat16",
+        "trust_remote_code": True,
+        "max_model_len": 2048,
+        "max_num_batched_tokens": 8192,
+        "enforce_eager": True,
+    }
+    stages = [
+        {
+            "stage_id": 0,
+            "stage_type": "llm",
+            "runtime": {"process": True, "devices": "0", "max_batch_size": 1},
+            "engine_args": {**common_args, "model_stage": "thinker", "gpu_memory_utilization": 0.1},
+        },
+        {
+            "stage_id": 1,
+            "stage_type": "llm",
+            "engine_input_source": [0],
+            "runtime": {"process": True, "devices": "1", "max_batch_size": 1, "connector_type": "queue"},
+            "engine_args": {**common_args, "model_stage": "talker", "gpu_memory_utilization": 0.1},
+        },
+    ]
+    connectors = [{"src_stage_id": 0, "dst_stage_id": 1, "connector_type": "queue"}]
+    engine = AsyncOmni(model=model_name, stages=stages, connectors=connectors, init_timeout=600, enable_sleep_mode=True)
+    yield engine
+    engine.shutdown()
+
+
+@pytest.fixture(scope="function")
+async def diffusion_engine():
+    if current_omni_platform.is_rocm():
+        clean_gpu_envs()
+    model_name = "ByteDance-Seed/BAGEL-7B-MoT"
+    stages = [
+        {
+            "stage_id": 0,
+            "stage_type": "diffusion",
+            "runtime": {"process": True, "devices": "0,1", "max_batch_size": 1},
+            "engine_args": {
+                "model_stage": "base",
+                "gpu_memory_utilization": 0.1,
+                "model_class_name": "BagelPipeline",
+                "enable_sleep_mode": True,
+                "enforce_eager": True,
+                "max_num_batched_tokens": 8192,
+                "parallel_config": {
+                    "tensor_parallel_size": 2,
+                },
+            },
+            "final_output": True,
+            "final_output_type": "image",
+        }
+    ]
+    engine = AsyncOmni(model=model_name, stages=stages, init_timeout=600, enable_sleep_mode=True)
+    yield engine
+    engine.shutdown()
+
+
+class TestOmniSleepMode:
+    @pytest.mark.asyncio
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+    async def test_llm_sleep_ack(self, llm_engine: AsyncOmni):
+        """LLM Thinker (GPU0) Signal and Physical Recycling Audit"""
+        try:
+            acks = await llm_engine.sleep(stage_ids=[0], level=2)
+            # Verification signal successful
+            assert all(get_ack_info(ack, "status") == "SUCCESS" for ack in acks)
+            # Verify physical recycling volume
+            total_freed_bytes = sum(get_ack_info(ack, "freed_bytes", 0) for ack in acks)
+            freed_gib = total_freed_bytes / 1024**3
+            logger.info(f"Thinker VRAM physically reclaimed: {freed_gib:.2f} GiB")
+            assert freed_gib > 5.0
+        finally:
+            llm_engine.shutdown()
+
+    @pytest.mark.asyncio
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
+    async def test_diffusion_sleep_handshake(self, diffusion_engine: AsyncOmni):
+        """Diffusion Worker stage signal loop"""
+        try:
+            logger.info("Starting Diffusion Worker Handshake Test")
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=2)
+
+            def _get_status(ack):
+                return ack.status if hasattr(ack, "status") else ack.get("status")
+
+            assert len(acks) >= 1, "Expected at least 1 ACK from Diffusion Workers"
+            assert all(_get_status(ack) == "SUCCESS" for ack in acks)
+            logger.info(f"Success: Received {len(acks)} Diffusion Worker ACKs")
+            logger.info("Testing auto-wakeup before test end...")
+            await diffusion_engine.wake_up(stage_ids=[0])
+            logger.info("Test logic finished, triggering manual shutdown...")
+        finally:
+            diffusion_engine.shutdown()
+            logger.info("Manual shutdown executed. Test should exit now.")
+
+    @pytest.mark.asyncio
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
+    async def test_cross_device_cleanup(self, diffusion_engine: AsyncOmni):
+        """Physical recycling audit: leveraging deterministic data returned by Workers"""
+        try:
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
+            # Sum up the release amounts reported by all Workers.
+            total_freed_bytes = sum(get_ack_info(ack, "freed_bytes", 0) for ack in acks)
+            freed_gb = total_freed_bytes / 1024**3
+            logger.info("Physical reclamation summary from workers:")
+            logger.info(f"- Total Workers: {len(acks)}")
+            logger.info(f"- Total Freed: {freed_gb:.2f} GiB")
+            assert freed_gb > 14.0
+            logger.info("SUCCESS: 100% weights offloaded.")
+        finally:
+            diffusion_engine.shutdown()
+
+    @pytest.mark.asyncio
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
+    async def test_diffusion_integrity_bit_level(self, diffusion_engine: AsyncOmni):
+        """Bit-level consistency after Diffusion wake-up (prevent image corruption)"""
+        try:
+            prompt = "A huge swimming pool, with many people swimming."
+            sp = OmniDiffusionSamplingParams(num_inference_steps=4, height=512, width=512, seed=42)
+            llm_sp = SamplingParams()
+
+            # Baseline Generation
+            logger.info("Running Baseline Generation...")
+            base_output = None
+            async for output in diffusion_engine.generate(prompt, request_id="base", sampling_params_list=[llm_sp, sp]):
+                base_output = output
+            assert base_output is not None and len(base_output.images) > 0
+            logger.info("Baseline Generation successful.")
+            # Sleep Level 2
+            logger.info("Entering Deep Sleep (VRAM Scavenging)...")
+            await diffusion_engine.sleep(stage_ids=[0], level=2)
+            # Wake-up
+            logger.info("Waking up (Reloading Weights)...")
+            await diffusion_engine.wake_up(stage_ids=[0])
+
+            await asyncio.sleep(2.0)
+            import gc
+
+            gc.collect()
+
+            logger.info("Running Post-Wakeup Generation...")
+            post_output = None
+            async for output in diffusion_engine.generate(prompt, request_id="post", sampling_params_list=[llm_sp, sp]):
+                post_output = output
+            # Assert result consistency
+            assert post_output is not None
+            assert len(base_output.images) == len(post_output.images)
+            assert post_output.images[0] is not None
+            logger.info("SUCCESS: Diffusion integrity verified after Sleep/Wake cycle.")
+        except Exception as e:
+            logger.error(f"Integrity test failed: {e}")
+            raise e
+        finally:
+            logger.info("Triggering mandatory cleanup...")
+            diffusion_engine.shutdown()
+            logger.info("Cleanup complete, test exiting.")
+
+    @pytest.mark.asyncio
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
+    async def test_coordinated_cross_device(self, llm_engine: AsyncOmni, diffusion_engine: AsyncOmni):
+        """Heterogeneous Coordinated Cleanup Test (Talker and Diffusion on GPU 1)"""
+        device_id = 1
+        try:
+            logger.info(f"Waking up both engines on GPU {device_id}...")
+            await llm_engine.wake_up(stage_ids=[1])
+            await diffusion_engine.wake_up(stage_ids=[0])
+
+            get_vram_info(device_id)
+            torch.cuda.empty_cache()
+            await asyncio.sleep(2)
+
+            initial_vram = get_vram_info(device_id)["reserved"]
+            logger.info(f"GPU {device_id} Peak Pressure: {initial_vram:.2f} GiB")
+
+            # coordinated sleep
+            logger.info("Issuing concurrent SLEEP commands...")
+            await llm_engine.sleep(stage_ids=[1], level=2)
+            await asyncio.sleep(1.0)
+            await diffusion_engine.sleep(stage_ids=[0], level=2)
+
+            await asyncio.sleep(3.0)
+            torch.cuda.empty_cache()
+
+            final_vram = get_vram_info(device_id)["reserved"]
+            logger.info(f"GPU {device_id} Final VRAM after coordinated sleep: {final_vram:.2f} GiB")
+
+            assert initial_vram - final_vram > 15.0 or final_vram < 8.0
+            logger.info(f"SUCCESS: Heterogeneous VRAM drop verified on GPU {device_id}.")
+        except Exception as e:
+            logger.error(f"Coordinated test failed: {e}")
+            raise e
+        finally:
+            logger.info("Triggering mandatory cleanup for both engines...")
+            llm_engine.shutdown()
+            diffusion_engine.shutdown()
+            logger.info("All engines scavenged. Ready for next test.")
+
+    @pytest.mark.asyncio
+    @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
+    async def test_diffusion_vram_lifecycle_audit(self, diffusion_engine: AsyncOmni):
+        """Diffusion memory loop: Active -> Deep Sleep -> Active -> inference sanity check"""
+        device_id = 1
+        try:
+            get_vram_info(device_id)
+            torch.cuda.empty_cache()
+            vram_initial = get_vram_info(device_id)["reserved"]
+            logger.info(f"Diffusion Initial VRAM: {vram_initial:.2f} GiB")
+
+            # Sleep
+            logger.info("Triggering Level 2 Deep Sleep (Full Weight Offloading)...")
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=2)
+
+            reported_freed_bytes = sum(getattr(ack, "freed_bytes", 0) for ack in acks)
+            reported_freed_gib = reported_freed_bytes / 1024**3
+            logger.info(f"Worker internally reported freed: {reported_freed_gib:.2f} GiB")
+
+            await asyncio.sleep(2)
+            get_vram_info(device_id)
+            torch.cuda.empty_cache()
+
+            vram_sleeping = get_vram_info(device_id)["reserved"]
+            logger.info(f"External VRAM measurement during Sleep: {vram_sleeping:.2f} GiB")
+
+            assert reported_freed_gib > 14.0 or vram_sleeping < 5.0, (
+                f"Reclamation failed. Reported: {reported_freed_gib:.2f}G, Measured: {vram_sleeping:.2f}G"
+            )
+
+            # wake-up
+            logger.info("Triggering Wake-up (Reloading weights to GPU)...")
+            await diffusion_engine.wake_up(stage_ids=[0])
+
+            await asyncio.sleep(2)
+            get_vram_info(device_id)
+            torch.cuda.empty_cache()
+            vram_restored = get_vram_info(device_id)["reserved"]
+            logger.info(f"VRAM after Wake-up: {vram_restored:.2f} GiB")
+
+            assert abs(vram_restored - vram_initial) < 3.0, "VRAM failed to restore to initial levels"
+
+            # inference sanity check
+            logger.info("Running post-lifecycle inference smoke test...")
+            prompt = "A futuristic lab with glowing lights, high quality."
+            sp = OmniDiffusionSamplingParams(num_inference_steps=2, height=512, width=512, seed=42)
+            llm_sp = SamplingParams()
+
+            base_img_found = False
+            async for output in diffusion_engine.generate(
+                prompt, request_id="lifecycle-check", sampling_params_list=[llm_sp, sp]
+            ):
+                if output.images and output.images[0] is not None:
+                    base_img_found = True
+
+            assert base_img_found, "Inference failed after Wake-up cycle!"
+            logger.info("SUCCESS: Full Diffusion Lifecycle (Active -> Sleep -> Active -> Generate) audited.")
+
+        except Exception as e:
+            logger.error(f"Lifecycle audit failed: {e}")
+            raise e
+        finally:
+            logger.info("Cleaning up engine and scavenging processes...")
+            diffusion_engine.shutdown()
+            await asyncio.sleep(1)

vllm_omni/config/model.py

@@ -112,6 +112,7 @@ class OmniModelConfig(ModelConfig):
     omni_kv_config: dict | None = None
     codec_frame_rate_hz: float | None = None
     task_type: str | None = None
+    enable_sleep_mode: bool = False
 
     @property
     def registry(self):

vllm_omni/diffusion/data.py

@@ -353,6 +353,8 @@ class DiffusionCacheConfig:
 @dataclass
 class OmniDiffusionConfig:
     # Model and path configuration (for convenience)
+    stage_id: int = 0
+
     model: str | None = None
 
     model_class_name: str | None = None
@@ -508,6 +510,8 @@ class OmniDiffusionConfig:
     # Step mode settings
     step_execution: bool = False
 
+    # sleep mode
+    enable_sleep_mode: bool = False
     # Maximum number of sequences to generate in a batch
     max_num_seqs: int = 1
 
@@ -797,5 +801,43 @@ class AttentionBackendEnum(enum.Enum):
         return self.name.lower()
 
 
+@dataclass
+class OmniACK:
+    """
+    Handshake payload from Workers to Orchestrator.
+    """
+
+    task_id: str
+    status: str
+    stage_id: int | None = None
+    rank: int | None = None
+    freed_bytes: int = 0
+    metadata: dict[str, Any] = field(default_factory=dict)
+    """
+    Additional telemetry such as:
+    - max_contiguous_block: for fragmentation analysis.
+    - cuda_graph_recalled: boolean if graphs were successfully destroyed/rebuilt.
+    - latency_ms: time taken for the D2H/H2D transfer.
+    """
+    error_msg: str | None = None
+
+
+@dataclass
+class OmniSleepTask:
+    """Structured sleep instruction."""
+
+    task_id: str
+    level: int = 2
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class OmniWakeTask:
+    """Structured wake-up instruction."""
+
+    task_id: str
+    tags: list[str] | None = None
+
+
 # Special message broadcast via scheduler queues to signal worker shutdown.
 SHUTDOWN_MESSAGE = {"type": "shutdown"}

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -36,10 +36,14 @@ class BackgroundResources:
 
     def __call__(self):
         """Clean up background resources."""
+        if hasattr(self, "wake_events") and self.wake_events:
+            for ev in self.wake_events:
+                ev.set()
+
         if self.broadcast_mq is not None:
             try:
                 for _ in range(self.num_workers):
-                    self.broadcast_mq.enqueue(SHUTDOWN_MESSAGE)
+                    self.broadcast_mq.enqueue(SHUTDOWN_MESSAGE, timeout=1.0)
 
                 self.broadcast_mq = None
                 self.result_mq = None
@@ -65,11 +69,13 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         self._closed = False
 
         num_workers = self.od_config.num_gpus
+        self.wake_events = [mp.Event() for _ in range(num_workers)]
+
         self._broadcast_mq = self._init_broadcast_queue(num_workers)
         broadcast_handle = self._broadcast_mq.export_handle()
 
         # Launch workers
-        processes, result_handle = self._launch_workers(broadcast_handle)
+        processes, result_handle = self._launch_workers(broadcast_handle, self.wake_events)
         self._result_mq = self._init_result_queue(result_handle)
         self._processes = processes
 
@@ -100,7 +106,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         if self._result_mq is None:
             raise RuntimeError("Result queue not initialized")
 
-    def _launch_workers(self, broadcast_handle):
+    def _launch_workers(self, broadcast_handle, wake_events):
         od_config = self.od_config
         logger.info("Starting server...")
 
@@ -126,6 +132,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                     od_config,
                     writer,
                     broadcast_handle,
+                    wake_events[i],
                     worker_extension_cls,
                     custom_pipeline_args,
                 ),

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -27,7 +27,10 @@ from vllm.v1.worker.workspace import init_workspace_manager
 
 from vllm_omni.diffusion.data import (
     DiffusionOutput,
+    OmniACK,
     OmniDiffusionConfig,
+    OmniSleepTask,
+    OmniWakeTask,
 )
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
@@ -77,6 +80,7 @@ class DiffusionWorker:
         self.model_runner: DiffusionModelRunner | None = None
         self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
         self.lora_manager: DiffusionLoRAManager | None = None
+        self.stage_id = getattr(od_config, "stage_id", 0)
         self.init_device()
         # Create model runner
         self.model_runner = DiffusionModelRunner(
@@ -159,8 +163,10 @@ class DiffusionWorker:
     def _get_profiler(self) -> WorkerProfiler | None:
         return getattr(self, "profiler", None)
 
-    def load_model(self, load_format: str = "default", custom_pipeline_name: str | None = None) -> None:
+    def load_model(self, load_format: str = "default", custom_pipeline_name: str | None = None, **kwargs) -> None:
         """Load the diffusion model using DiffusionModelRunner."""
+        load_format = kwargs.get("load_format", load_format)
+        custom_pipeline_name = kwargs.get("custom_pipeline_name", custom_pipeline_name)
         with (
             set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config),
             set_current_vllm_config(self.vllm_config),
@@ -170,6 +176,8 @@ class DiffusionWorker:
                 load_format=load_format,
                 custom_pipeline_name=custom_pipeline_name,
             )
+            current_omni_platform.synchronize()
+            gc.collect()
         process_memory = get_process_gpu_memory(self.local_rank)
         if process_memory is not None:
             logger.info(
@@ -284,77 +292,176 @@ class DiffusionWorker:
         """
         from vllm.device_allocator.cumem import CuMemAllocator
 
-        process_memory_before_sleep = get_process_gpu_memory(self.local_rank)
-        free_bytes_before_sleep = None
-        if process_memory_before_sleep is None:
-            free_bytes_before_sleep = current_omni_platform.get_free_memory()
+        allocator = CuMemAllocator.get_instance()
+
+        usage_before = allocator.get_current_usage()
 
-        # Save the buffers before level 2 sleep
         if level == 2 and self.model_runner is not None:
+            if hasattr(self.model_runner, "graph_runners"):
+                self.model_runner.graph_runners.clear()
+                logger.info(f"[Worker {self.rank}] CUDA Graphs cleared.")
             model = self.model_runner.pipeline
             self._sleep_saved_buffers = {name: buffer.cpu().clone() for name, buffer in model.named_buffers()}
 
-        allocator = CuMemAllocator.get_instance()
-        allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
-        process_memory_after_sleep = get_process_gpu_memory(self.local_rank)
-        if process_memory_before_sleep is not None and process_memory_after_sleep is not None:
-            freed_bytes = process_memory_before_sleep - process_memory_after_sleep
-            used_bytes = process_memory_after_sleep
-            accounting_scope = "process-scoped"
+        free_mem_before = current_omni_platform.get_free_memory()
+
+        # Level 1: Offload weights; Level 2: Total Discard
+        offload_tags = ("weights",) if level == 1 else tuple()
+        allocator.sleep(offload_tags=offload_tags)
+
+        current_omni_platform.empty_cache()
+        current_omni_platform.synchronize()
+
+        free_mem_after = current_omni_platform.get_free_memory()
+        try:
+            total_mem = current_omni_platform.get_device_total_memory()
+        except (NotImplementedError, AttributeError):
+            total_mem = torch.cuda.get_device_properties(self.device).total_memory
+
+        phys_freed_bytes = max(0, free_mem_after - free_mem_before)
+        phys_used_bytes = total_mem - free_mem_after
+
+        if usage_before > 0:
+            logger.info(
+                f"[Diffusion Worker {self.rank}] Sleep Level {level}: "
+                f"physically freed {phys_freed_bytes / GiB_bytes:.2f} GiB, "
+                f"{phys_used_bytes / GiB_bytes:.2f} GiB is still in use."
+            )
         else:
-            free_bytes_after_sleep = current_omni_platform.get_free_memory()
-            assert free_bytes_before_sleep is not None
-            device_id = self.device.index if self.device.index is not None else 0
-            total = current_omni_platform.get_device_total_memory(device_id)
-            freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
-            used_bytes = total - free_bytes_after_sleep
-            accounting_scope = "device-scoped fallback"
-        assert freed_bytes >= 0, "Memory usage increased after sleeping."
-        logger.info(
-            "Sleep mode (%s) freed %.2f GiB memory, %.2f GiB memory is still in use.",
-            accounting_scope,
-            freed_bytes / GiB_bytes,
-            used_bytes / GiB_bytes,
-        )
-        return True
+            logger.info(f"[Worker {self.rank}] Sleep Level {level} completed (GPU was already empty).")
+        logger.info(f"[Worker {self.rank}] Memory usage before sleep: {usage_before / GiB_bytes:.2f} GiB.")
+        return usage_before
 
     def wake_up(self, tags: list[str] | None = None) -> bool:
         """
-        Wake up the worker from sleep mode. See the sleep function
-        method for more details.
+        Wake up the worker from sleep mode.
+
+        Re-activates the memory allocator for the specified tags and restores
+        model buffers from CPU back to GPU if they were saved during Level 2 sleep.
 
         Args:
-            tags: An optional list of tags to reallocate the worker memory
-                for specific memory allocations. Values must be in
-                `("weights")`. If None, all memory is reallocated.
-                wake_up should be called with all tags (or None) before the
-                worker is used again.
+            tags: List of memory pool tags to re-activate (e.g., ["weights"]
+                  to match Level 1 sleep). If None, all pools are re-activated.
         """
         from vllm.device_allocator.cumem import CuMemAllocator
 
         allocator = CuMemAllocator.get_instance()
         allocator.wake_up(tags)
-
-        # Restore the buffers after level 2 sleep
+        current_omni_platform.synchronize()
         if len(self._sleep_saved_buffers) and self.model_runner is not None:
             model = self.model_runner.pipeline
             for name, buffer in model.named_buffers():
                 if name in self._sleep_saved_buffers:
                     buffer.data.copy_(self._sleep_saved_buffers[name].data)
             self._sleep_saved_buffers = {}
+            logger.info(f"[Worker {self.rank}] Buffers restored from CPU.")
+        logger.info(f"[Worker {self.rank}] Wake-up complete.")
         return True
 
+    def handle_sleep_task(self, task: OmniSleepTask) -> OmniACK:
+        from vllm_omni.platforms import current_omni_platform
+
+        try:
+            if isinstance(task, dict):
+                task = OmniSleepTask(**task)
+            logger.info(f"[Worker {self.rank}] Handshake Received: Task {task.task_id}")
+
+            current_omni_platform.synchronize()
+            usage_before = current_omni_platform.get_current_memory_usage(self.device)
+            self.sleep(level=task.level)
+            current_omni_platform.synchronize()
+            usage_after = current_omni_platform.get_current_memory_usage(self.device)
+            real_freed = max(0, usage_before - usage_after)
+            logger.info(f"[Worker {self.rank}] Preparing ACK: freed_bytes={real_freed / GiB_bytes:.2f} GiB.")
+
+            # Ensure all ranks have completed sleep before measuring memory and sending ACK
+            if torch.distributed.is_initialized():
+                t_freed = torch.tensor([float(real_freed)], device=self.device)
+                torch.distributed.all_reduce(t_freed)
+                real_freed = int(t_freed.item())
+
+            if self.rank != 0:
+                return None
+
+            ack = OmniACK(
+                task_id=task.task_id,
+                status="SUCCESS",
+                stage_id=self.stage_id,
+                rank=self.rank,
+                freed_bytes=real_freed,
+                # return RL need metadata
+                metadata={
+                    "source": f"Platform_{current_omni_platform.get_device_name()}",
+                    "total_freed_gib": f"{real_freed / GiB_bytes:.2f}",
+                    "rank_residual_gib": f"{usage_after / GiB_bytes:.2f}",
+                },
+            )
+            logger.info(f"[Worker {self.rank}] ACK emitted. Freed {real_freed / GiB_bytes:.2f} GiB.")
+            return ack
+        except Exception as e:
+            logger.error(f"Sleep failed: {e}", exc_info=True)
+            if torch.distributed.is_initialized():
+                try:
+                    torch.distributed.barrier()
+                except Exception:
+                    pass
+            return OmniACK(task_id=task.task_id, status="ERROR", error_msg=str(e))
+
+    def handle_wake_task(self, task: OmniWakeTask) -> OmniACK:
+        from vllm_omni.platforms import current_omni_platform
+
+        try:
+            if isinstance(task, dict):
+                task = OmniWakeTask(**task)
+            logger.info(f"[Worker {self.rank}] Responding to Wake-up Task: {task.task_id}")
+            self.wake_up(tags=task.tags)
+
+            logger.info(f"[Worker {self.rank}] wake_up logic finished, entering barrier...")
+            if torch.distributed.is_initialized():
+                torch.distributed.barrier()
+
+            current_omni_platform.synchronize()
+            usage_now = current_omni_platform.get_current_memory_usage(self.device)
+            current_used_gib = usage_now / (1024**3)
+
+            if self.rank != 0:
+                return None
+            logger.info(f"[Worker {self.rank}] PASSED barrier, about to return to loop.")
+
+            return OmniACK(
+                task_id=task.task_id,
+                status="SUCCESS",
+                stage_id=self.stage_id,
+                rank=self.rank,
+                metadata={
+                    "state": "WARM",
+                    "source": f"Platform_{current_omni_platform.get_device_name()}",
+                    "current_vram_gib": f"{current_used_gib:.2f}",
+                },
+            )
+        except Exception as e:
+            logger.error(f"Wake-up failed on Rank {self.rank}: {e}", exc_info=True)
+            if torch.distributed.is_initialized():
+                try:
+                    torch.distributed.barrier()
+                except Exception:
+                    pass
+            return OmniACK(task_id=task.task_id, status="ERROR", error_msg=str(e))
+
     def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
         """Get memory pool context for sleep mode support."""
-        if self.od_config.enable_sleep_mode:
+        is_sleep_enabled = getattr(self.od_config, "enable_sleep_mode", False)
+        if is_sleep_enabled:
+            current_omni_platform.synchronize()
+            gc.collect()
             from vllm.device_allocator.cumem import CuMemAllocator
 
             allocator = CuMemAllocator.get_instance()
             if tag == "weights":
                 assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process."
+            logger.info(f"[Worker {self.rank}] Activating Diffusion CuMem pool for tag: {tag}")
             return allocator.use_memory_pool(tag=tag)
-        else:
-            return nullcontext()
+        return nullcontext()
 
     def shutdown(self) -> None:
         """Shutdown the worker and cleanup distributed environment."""
@@ -395,11 +502,13 @@ class WorkerProc:
         od_config: OmniDiffusionConfig,
         gpu_id: int,
         broadcast_handle,
+        wake_event: mp.Event,
         worker_extension_cls: str | None = None,
         custom_pipeline_args: dict[str, Any] | None = None,
     ):
         self.od_config = od_config
         self.gpu_id = gpu_id
+        self.wake_event = wake_event
 
         # Inter-process Communication
         self.context = zmq.Context(io_threads=2)
@@ -414,7 +523,13 @@ class WorkerProc:
         if gpu_id == 0:
             self.result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0])
             self.result_mq_handle = self.result_mq.export_handle()
+            WorkerProc._shared_result_handle = self.result_mq_handle
             logger.info(f"Worker {gpu_id} created result MessageQueue")
+        else:
+            handle = getattr(WorkerProc, "_shared_result_handle", None)
+            if handle:
+                self.result_mq = MessageQueue.create_from_handle(handle, gpu_id)
+                logger.info(f"Worker {gpu_id} attached to shared result MessageQueue")
 
         assert od_config.master_port is not None
 
@@ -438,13 +553,17 @@ class WorkerProc:
         )
         return wrapper
 
-    def return_result(self, output: object):
+    def return_result(self, output: Any):
         """Reply to client, only on rank 0."""
         if self.result_mq is not None:
+            if isinstance(output, OmniACK):
+                self.result_mq.enqueue(output)
+                return
             try:
                 pack_diffusion_output_shm(output)
             except Exception as e:
-                logger.warning("SHM pack failed, falling back to raw enqueue: %s", e)
+                if hasattr(output, "output"):
+                    logger.warning("SHM pack failed for model output: %s", e)
             self.result_mq.enqueue(output)
 
     def recv_message(self):
@@ -480,20 +599,31 @@ class WorkerProc:
         while self._running:
             msg = None
             try:
-                msg = self.recv_message()
-            except Exception as e:
-                logger.error(
-                    f"Error receiving message in worker loop: {e}",
-                    exc_info=True,
-                )
+                msg = self.mq.dequeue(timeout=1.0)
+            except Exception:
+                if self.wake_event and self.wake_event.is_set():
+                    self.wake_event.clear()
+                    logger.info(f"Worker {self.gpu_id} caught OOB POKE, forcing wake-up sequence.")
+                    msg = {"type": "wake_up", "task_id": "recovery-task", "tags": None}
+                else:
+                    continue
+            if msg is None:
                 continue
 
             if msg is None or len(msg) == 0:
                 logger.warning("Worker %s: Received empty payload, ignoring", self.gpu_id)
                 continue
 
+            if isinstance(msg, dict) and msg.get("type") == "sleep":
+                task = OmniSleepTask(level=msg.get("level", 2), task_id=msg.get("task_id", "local"))
+                ack = self.worker.handle_sleep_task(task)
+                self.return_result(ack)
+            elif isinstance(msg, dict) and msg.get("type") == "wake_up":
+                task = OmniWakeTask(tags=msg.get("tags"), task_id=msg.get("task_id", "local"))
+                ack = self.worker.handle_wake_task(task)
+                self.return_result(ack)
             # Route message based on type
-            if isinstance(msg, dict) and msg.get("type") == "rpc":
+            elif isinstance(msg, dict) and msg.get("type") == "rpc":
                 try:
                     result, should_reply = self.execute_rpc(msg)
                     if should_reply:
@@ -538,6 +668,7 @@ class WorkerProc:
         od_config: OmniDiffusionConfig,
         pipe_writer: mp.connection.Connection,
         broadcast_handle,
+        wake_event: mp.Event,
         worker_extension_cls: str | None = None,
         custom_pipeline_args: dict[str, Any] | None = None,
     ) -> None:
@@ -549,6 +680,7 @@ class WorkerProc:
             od_config,
             gpu_id=rank,
             broadcast_handle=broadcast_handle,
+            wake_event=wake_event,
             worker_extension_cls=worker_extension_cls,
             custom_pipeline_args=custom_pipeline_args,
         )
@@ -574,6 +706,7 @@ class WorkerWrapperBase:
         gpu_id: int,
         od_config: OmniDiffusionConfig,
         base_worker_class: type = DiffusionWorker,
+        wake_event: mp.Event = None,
         worker_extension_cls: str | None = None,
         custom_pipeline_args: dict[str, Any] | None = None,
     ):
@@ -734,6 +867,12 @@ class WorkerWrapperBase:
         """
         return self.worker.wake_up(tags)
 
+    def handle_sleep_task(self, task):
+        return self.worker.handle_sleep_task(task)
+
+    def handle_wake_task(self, task):
+        return self.worker.handle_wake_task(task)
+
     def shutdown(self) -> None:
         """Shutdown the worker and cleanup resources."""
         return self.worker.shutdown()

vllm_omni/engine/arg_utils.py

@@ -6,11 +6,12 @@ import tempfile
 from dataclasses import dataclass, field, fields
 from typing import Any
 
-from vllm.engine.arg_utils import EngineArgs
+from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
 from vllm.logger import init_logger
 
 from vllm_omni.config import OmniModelConfig
 from vllm_omni.engine.output_modality import OutputModality
+from vllm_omni.platforms import current_omni_platform
 from vllm_omni.plugins import load_omni_general_plugins
 
 logger = init_logger(__name__)
@@ -144,6 +145,24 @@ class OmniEngineArgs(EngineArgs):
     quantization_config: Any | None = None
     worker_type: str | None = None
     task_type: str | None = None
+    worker_cls: str = None
+    enable_sleep_mode: bool = False
+    omni: bool = False
+
+    @classmethod
+    def _add_omni_specific_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+        try:
+            parser.add_argument("--omni", action="store_true", default=False, help="Enable Omni engine features.")
+        except argparse.ArgumentError:
+            pass
+        try:
+            parser.add_argument(
+                "--enable-sleep-mode", action="store_true", default=False, help="Enable GPU memory pool for sleep mode."
+            )
+        except argparse.ArgumentError:
+            pass
+        return parser
+
     omni_master_address: str | None = None
     omni_master_port: int | None = None
     stage_configs_path: str | None = None
@@ -152,6 +171,11 @@ class OmniEngineArgs(EngineArgs):
     custom_pipeline_args: dict[str, Any] | None = None
 
     def __post_init__(self) -> None:
+        if self.worker_cls is None:
+            if self.worker_type == "ar":
+                self.worker_cls = current_omni_platform.get_omni_ar_worker_cls()
+            elif self.worker_type == "generation":
+                self.worker_cls = current_omni_platform.get_omni_generation_worker_cls()
         load_omni_general_plugins()
         super().__post_init__()
 
@@ -296,6 +320,15 @@ class OmniEngineArgs(EngineArgs):
         )
         return omni_config
 
+
+@dataclass
+class OmniAsyncEngineArgs(AsyncEngineArgs, OmniEngineArgs):
+    @classmethod
+    def add_cli_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+        parser = AsyncEngineArgs.add_cli_args(parser)
+        parser = OmniEngineArgs._add_omni_specific_args(parser)
+        return parser
+
     @property
     def output_modality(self) -> OutputModality:
         """Parse engine_output_type into a type-safe OutputModality flag."""

vllm_omni/engine/async_omni_engine.py

@@ -793,6 +793,7 @@ class AsyncOmniEngine:
                                 else:
                                     use_inline = True if self.num_stages == 1 else False
                                     stage_clients[stage_idx] = initialize_diffusion_stage(
+                                        configured_stage_id,
                                         self.model,
                                         stage_cfg,
                                         metadata,
@@ -1401,6 +1402,12 @@ class AsyncOmniEngine:
         # Inject diffusion LoRA-related knobs from kwargs if not present in the stage config.
         for cfg in stage_configs:
             try:
+                if not hasattr(cfg, "engine_args") or cfg.engine_args is None:
+                    cfg.engine_args = OmegaConf.create({})
+                global_sleep_mode = kwargs.get("enable_sleep_mode")
+                if global_sleep_mode is not None:
+                    if not hasattr(cfg.engine_args, "enable_sleep_mode") or cfg.engine_args.enable_sleep_mode is None:
+                        cfg.engine_args.enable_sleep_mode = global_sleep_mode
                 if getattr(cfg, "stage_type", None) != "diffusion":
                     continue
                 if not hasattr(cfg, "engine_args") or cfg.engine_args is None:

vllm_omni/engine/stage_init_utils.py

@@ -23,6 +23,7 @@ from vllm.usage.usage_lib import UsageContext
 from vllm.v1.engine.input_processor import InputProcessor
 from vllm.v1.executor import Executor
 
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices
 from vllm_omni.entrypoints.utils import filter_dataclass_kwargs, resolve_model_config_path
@@ -650,6 +651,7 @@ def build_diffusion_config(
 
 
 def initialize_diffusion_stage(
+    stage_id: int,
     model: str,
     stage_cfg: Any,
     metadata: StageMetadata,
@@ -671,6 +673,15 @@ def initialize_diffusion_stage(
     """
     from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client
 
+    engine_args = _to_dict(stage_cfg.engine_args)
+    engine_args.pop("stage_id", None)
+    od_config = OmniDiffusionConfig.from_kwargs(
+        stage_id=stage_id,
+        model=model,
+        **engine_args,
+    )
+    if metadata.cfg_kv_collect_func is not None:
+        od_config.cfg_kv_collect_func = metadata.cfg_kv_collect_func
     od_config = build_diffusion_config(model, stage_cfg, metadata)
     return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline)
 

vllm_omni/entrypoints/async_omni.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 
 import asyncio
 import time
+import uuid
 from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence
 from typing import TYPE_CHECKING, Any
 
@@ -24,10 +25,12 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.tasks import SupportedTask
 from vllm.v1.engine.exceptions import EngineDeadError
 
+from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni_base import OmniBase
 from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
 
 if TYPE_CHECKING:
     from vllm.inputs.preprocess import InputPreprocessor
@@ -40,6 +43,63 @@ logger = init_logger(__name__)
 _FINAL_OUTPUT_IDLE_SLEEP_S = 0.001
 
 
+class AsyncEventResolver:
+    """
+    A generic signal aggregator designed for synchronized handshakes in
+    distributed or multi-stage environments. Supports waiting for a specified
+    number (expected_count) of worker signals in both inline and multiprocess modes.
+    """
+
+    def __init__(self, orchestrator=None):
+        self._pending_tasks: dict[str, dict] = {}
+        self.orchestrator = orchestrator
+        self._lock = asyncio.Lock()
+
+    def watch_task(self, task_id: str, expected_count: int = 1) -> asyncio.Future:
+        loop = asyncio.get_running_loop()
+        fut = loop.create_future()
+        self._pending_tasks[task_id] = {
+            "future": fut,
+            "expected_count": expected_count,
+            "received": [],
+            "start_time": time.time(),
+        }
+        return fut
+
+    async def resolve(self, ack: OmniACK):
+        tid = getattr(ack, "task_id", None)
+
+        if tid is None and isinstance(ack, dict):
+            tid = ack.get("task_id")
+
+        async with self._lock:
+            task_info = self._pending_tasks.get(tid)
+            if task_info is None:
+                logger.warning(f"Received stray ACK for task_id {tid}. Task might have timed out.")
+                return
+
+            task_info["received"].append(ack)
+            current_count = len(task_info["received"])
+            expected = task_info["expected_count"]
+
+            orchestrator = self.orchestrator
+            if orchestrator and hasattr(orchestrator, "metrics") and orchestrator.metrics:
+                freed = getattr(ack, "freed_bytes", 0)
+                if freed == 0 and isinstance(ack, dict):
+                    freed = ack.get("freed_bytes", 0)
+                orchestrator.metrics.record_vram_reclaimed(freed)
+
+            logger.info(f"[Resolver] Task {tid} progress: {current_count}/{expected} ACKs received.")
+
+            if current_count >= expected:
+                self._pending_tasks.pop(tid)
+                fut = task_info["future"]
+                if not fut.done():
+                    elapsed = time.time() - task_info["start_time"]
+                    logger.info(f"[Resolver] Task {tid} completed successfully in {elapsed:.2f}s.")
+                    fut.set_result(task_info["received"])
+
+
 class AsyncOmni(EngineClient, OmniBase):
     """Asynchronous unified entry point for multi-stage pipelines using AsyncOmniEngine.
 
@@ -76,7 +136,7 @@ class AsyncOmni(EngineClient, OmniBase):
         self._paused: bool = False
         self._is_sleeping: bool = False
         self.final_output_task: asyncio.Task | None = None
-
+        self.event_resolver = AsyncEventResolver(orchestrator=self)
         self.config_path = self.engine.config_path
         self.tts_max_instructions_length = kwargs.get("tts_max_instructions_length", None)
         self.input_processor = self.engine.input_processor
@@ -488,6 +548,22 @@ class AsyncOmni(EngineClient, OmniBase):
                         await asyncio.sleep(_FINAL_OUTPUT_IDLE_SLEEP_S)
                         continue
 
+                    if isinstance(msg, dict) and msg.get("type") == "ack":
+                        ack_data = msg.get("ack")
+                        tid = getattr(ack_data, "task_id", "unknown")
+                        logger.info(f"[{self._name}] Intercepted wrapped ACK for task {tid}")
+                        await self.event_resolver.resolve(ack_data)
+                        continue
+                    if isinstance(msg, OmniACK):
+                        logger.info(f"[{self._name}] Intercepted raw ACK object: {msg.task_id}")
+                        await self.event_resolver.resolve(msg)
+                        continue
+                    if hasattr(msg, "task_id"):
+                        tid = getattr(msg, "task_id")
+                        logger.info(f"[{self._name}] Intercepted task-ID object: {tid}")
+                        await self.event_resolver.resolve(msg)
+                        continue
+
                     should_continue, _, stage_id, req_state = self._handle_output_message(msg)
                     if should_continue:
                         continue
@@ -654,21 +730,68 @@ class AsyncOmni(EngineClient, OmniBase):
         logger.warning("[AsyncOmni] reset_prefix_cache not yet supported with Orchestrator process")
         return True
 
-    async def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
-        """Sleep all stages.
+    async def sleep(
+        self, stage_ids: list[int] | None = None, level: int = 2, mode: PauseMode = "abort"
+    ) -> list[OmniACK]:
+        self._final_output_handler()
+        if stage_ids is None:
+            stage_ids = list(range(len(self.engine.stage_clients)))
+        total_workers = 0
+        for sid in stage_ids:
+            client = self.engine.stage_clients[sid]
+            # During the Diffusion phase, regardless of the TP amount,
+            # currently only a summary ACK is reported at Rank 0.
+            if getattr(client, "stage_type", "") == "diffusion":
+                total_workers += 1
+            else:
+                config = self.engine.stage_vllm_configs[sid]
+                actual_tp = config.parallel_config.tensor_parallel_size if config else 1
+                total_workers += actual_tp
 
-        Best-effort: unsupported stages will emit a TODO result.
-        """
+        task_id = str(uuid.uuid4())
+        self.event_resolver.watch_task(task_id, expected_count=total_workers)
+        logger.info(f"[{self._name}] Sleep initiated (Task: {task_id}). Awaiting {total_workers} ACKs...")
+        task = OmniSleepTask(level=level, task_id=task_id)
+        rpc_results = await self.collective_rpc(method="handle_sleep_task", args=(task,), stage_ids=stage_ids)
+        final_acks = []
+        for stage_res in rpc_results:
+            worker_acks = stage_res if isinstance(stage_res, list) else [stage_res]
+            for ack in worker_acks:
+                if ack is not None:
+                    await self.event_resolver.resolve(ack)
+                    final_acks.append(ack)
         self._is_sleeping = True
-        await self.collective_rpc(method="sleep", args=(level,))
+        return final_acks
 
-    async def wake_up(self, tags: list[str] | None = None) -> None:
-        """Wake up all stages.
-
-        Best-effort: unsupported stages will emit a TODO result.
-        """
+    async def wake_up(self, stage_ids: list[int] | None = None, tags: list[str] | None = None) -> list[OmniACK]:
+        self._final_output_handler()
+        if stage_ids is None:
+            stage_ids = list(range(len(self.engine.stage_clients)))
+        total_workers = 0
+        for sid in stage_ids:
+            client = self.engine.stage_clients[sid]
+            if getattr(client, "stage_type", "") == "diffusion":
+                total_workers += 1
+            else:
+                config = self.engine.stage_vllm_configs[sid]
+                total_workers += config.parallel_config.tensor_parallel_size if config else 1
+        task_id = str(uuid.uuid4())
+        self.event_resolver.watch_task(task_id, expected_count=total_workers)
+        logger.info(f"[{self._name}] Wake-up initiated (Task: {task_id}). Awaiting {total_workers} ACKs...")
+        task = OmniWakeTask(tags=tags, task_id=task_id)
+        rpc_results = await self.collective_rpc(method="handle_wake_task", args=(task,), stage_ids=stage_ids)
+        final_acks = []
+        for stage_res in rpc_results:
+            worker_acks = stage_res if isinstance(stage_res, list) else [stage_res]
+            for ack in worker_acks:
+                if ack is not None:
+                    await self.event_resolver.resolve(ack)
+                    final_acks.append(ack)
+        current_omni_platform.synchronize()
+        await asyncio.sleep(0.1)
         self._is_sleeping = False
-        await self.collective_rpc(method="wake_up", args=(tags,))
+        logger.info(f"[{self._name}] All {len(final_acks)}/{total_workers} workers reported WARM for task {task_id}.")
+        return final_acks
 
     async def is_sleeping(self) -> bool:
         """Return whether all stages are sleeping.
@@ -725,6 +848,10 @@ class AsyncOmni(EngineClient, OmniBase):
         """Whether orchestrator thread has stopped unexpectedly."""
         return not self.engine.is_alive()
 
+    @property
+    def _name(self) -> str:
+        return "AsyncOrchestrator"
+
     @property
     def is_stopped(self) -> bool:
         """EngineClient abstract property implementation."""

vllm_omni/entrypoints/cli/serve.py

@@ -139,6 +139,17 @@ class OmniServeCommand(CLISubcommand):
             action="store_true",
             help="Enable vLLM-Omni mode for multi-modal and diffusion models",
         )
+
+        try:
+            omni_config_group.add_argument(
+                "--enable-sleep-mode",
+                action="store_true",
+                default=False,
+                help="Enable GPU memory pool for sleep mode.",
+            )
+        except argparse.ArgumentError:
+            pass
+
         omni_config_group.add_argument(
             "--task-type",
             type=str,

vllm_omni/entrypoints/omni_base.py

@@ -133,7 +133,7 @@ class OmniBase(PDDisaggregationMixin):
         if "log_requests" in kwargs:
             raise TypeError("`log_requests` has been removed in Omni/AsyncOmni. Use `log_stats`.")
         model = omni_snapshot_download(model)
-        self._name = self.__class__.__name__
+        self.__dict__["_name"] = self.__class__.__name__
         self.model = model
         self.log_stats = log_stats
         # Provisional value (mirrors the CLI/caller kwarg); the engine resolves

vllm_omni/entrypoints/openai/api_server.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import asyncio
 import base64
+import dataclasses
 import io
 import json
 import multiprocessing
@@ -837,6 +838,7 @@ async def omni_init_app_state(
 
     state.enable_server_load_tracking = args.enable_server_load_tracking
     state.server_load_metrics = 0
+    state.sleeping_stages = set()
 
 
 def Omnivideo(request: Request) -> OmniOpenAIServingVideo | None:
@@ -2553,3 +2555,64 @@ async def stop_profile(raw_request: Request, request: ProfileRequest | None = No
         raise HTTPException(
             status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Failed to stop profiler: {str(e)}"
         )
+
+
+class OmniSleepRequest(BaseModel):
+    stage_ids: list[int]
+    level: int = 2
+
+
+class OmniWakeupRequest(BaseModel):
+    stage_ids: list[int]
+
+
+@router.post("/v1/omni/sleep")
+async def omni_sleep(request: OmniSleepRequest, raw_request: Request):
+    engine_client = raw_request.app.state.engine_client
+    sleeping_set = raw_request.app.state.sleeping_stages
+    if not hasattr(engine_client, "sleep"):
+        raise HTTPException(status_code=501, detail="Engine does not support sleep")
+    acks = await engine_client.sleep(stage_ids=request.stage_ids, level=request.level)
+    for sid in request.stage_ids:
+        sleeping_set.add(sid)
+    return {"status": "SUCCESS", "acks": [dataclasses.asdict(a) if dataclasses.is_dataclass(a) else a for a in acks]}
+
+
+@router.post("/v1/omni/wakeup")
+async def omni_wakeup(request: OmniWakeupRequest, raw_request: Request):
+    engine_client = raw_request.app.state.engine_client
+    sleeping_set = raw_request.app.state.sleeping_stages
+    if not any(sid in sleeping_set for sid in request.stage_ids):
+        return {"status": "SKIPPED", "reason": "Target stages are not sleeping."}
+    if not hasattr(engine_client, "wake_up"):
+        raise HTTPException(status_code=501, detail="Engine does not support wake_up")
+    acks = await engine_client.wake_up(stage_ids=request.stage_ids)
+    for sid in request.stage_ids:
+        if sid in sleeping_set:
+            sleeping_set.remove(sid)
+    return {"status": "SUCCESS", "acks": [dataclasses.asdict(a) if dataclasses.is_dataclass(a) else a for a in acks]}
+
+
+if __name__ == "__main__":
+    import argparse
+    import asyncio
+
+    from vllm.entrypoints.openai.cli_args import make_arg_parser
+
+    parser = argparse.ArgumentParser(description="vLLM-Omni OpenAI-Compatible REST API server")
+    parser = make_arg_parser(parser)
+    registered_flags = set()
+    for action in parser._actions:
+        registered_flags.update(action.option_strings)
+    if "--omni" not in registered_flags:
+        parser.add_argument("--omni", action="store_true", default=False, help="Enable vLLM-Omni mode.")
+    if "--enable-sleep-mode" not in registered_flags:
+        parser.add_argument(
+            "--enable-sleep-mode", action="store_true", default=False, help="Enable GPU memory pool for sleep mode."
+        )
+    args = parser.parse_args()
+    if not hasattr(args, "model_tag"):
+        setattr(args, "model_tag", args.model)
+    if hasattr(args, "model_tag") and args.model_tag is None:
+        args.model_tag = args.model
+    asyncio.run(omni_run_server(args))

vllm_omni/worker/base.py

@@ -2,15 +2,23 @@
 
 from __future__ import annotations
 
+import gc
 import os
 import time
+from contextlib import AbstractContextManager, nullcontext
 
 import torch
 from vllm.logger import init_logger
 from vllm.utils.mem_utils import format_gib, memory_profiling
 from vllm.v1.worker.gpu_worker import Worker as GPUWorker
 
+from vllm_omni.diffusion.data import (
+    OmniACK,
+    OmniSleepTask,
+    OmniWakeTask,
+)
 from vllm_omni.entrypoints.utils import detect_pid_host
+from vllm_omni.platforms import current_omni_platform
 from vllm_omni.worker.gpu_memory_utils import (
     get_process_gpu_memory,
     is_process_scoped_memory_available,
@@ -30,6 +38,13 @@ class OmniGPUWorkerBase(GPUWorker):
     for custom trace naming, background gzip, and trace path collection.
     """
 
+    def load_model(self, *args, **kwargs):
+        with self._maybe_get_memory_pool_context("weights"):
+            res = super().load_model(*args, **kwargs)
+            current_omni_platform.synchronize()
+            gc.collect()
+            return res
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -94,10 +109,8 @@ class OmniGPUWorkerBase(GPUWorker):
         """
         if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes:
             self.model_runner.profile_run()
-            logger.info(
-                "Using explicit kv_cache_memory_bytes: %s GiB",
-                format_gib(kv_cache_memory_bytes),
-            )
+            if current_omni_platform.is_rocm():
+                torch.cuda.synchronize()
             return kv_cache_memory_bytes
 
         with memory_profiling(
@@ -154,3 +167,141 @@ class OmniGPUWorkerBase(GPUWorker):
             )
 
         return int(self.available_kv_cache_memory_bytes)
+
+    # Provide memory pool context
+    def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
+        v1_config_enabled = False
+        if hasattr(self, "vllm_config"):
+            model_cfg = getattr(self.vllm_config, "model_config", None)
+            v1_config_enabled = getattr(model_cfg, "enable_sleep_mode", False)
+
+        is_sleep_enabled = v1_config_enabled or getattr(self.cache_config, "enable_sleep_mode", False)
+        if is_sleep_enabled:
+            current_omni_platform.synchronize()
+            gc.collect()
+            from vllm.device_allocator.cumem import CuMemAllocator
+
+            allocator = CuMemAllocator.get_instance()
+            logger.info(f"[LLM Worker {self.rank}] Sleep Mode ENABLED. Activating CuMem pool for tag: {tag}")
+            return allocator.use_memory_pool(tag=tag)
+        else:
+            logger.warning(f"[LLM Worker {self.rank}] Sleep Mode DISABLED.")
+            return nullcontext()
+
+    def sleep(self, level: int = 1) -> bool:
+        """
+        Put the worker to sleep.
+        Args:
+            level: 1 (Offload weights to CPU), level: 2 (Total Discard).
+        """
+        from vllm.device_allocator.cumem import CuMemAllocator
+
+        mem_before = current_omni_platform.get_current_memory_usage(self.device)
+        offload_tags = ("weights",) if level == 1 else tuple()
+        allocator = CuMemAllocator.get_instance()
+        allocator.sleep(offload_tags=offload_tags)
+        current_omni_platform.empty_cache()
+        current_omni_platform.synchronize()
+        mem_after = current_omni_platform.get_current_memory_usage(self.device)
+        freed = max(0, mem_before - mem_after)
+        remaining_gb = mem_after / 1024**3
+        logger.info(
+            f"[LLM Worker {self.rank}] Level {level} Sleep: Freed "
+            f"{freed / 1024**3:.2f} GiB. {remaining_gb:.2f}GiB memory "
+            "is still in use."
+        )
+        return True
+
+    def wake_up(self, tags: list[str] | None = None) -> bool:
+        "Physical video memory reloading logic"
+        from vllm.device_allocator.cumem import CuMemAllocator
+
+        allocator = CuMemAllocator.get_instance()
+        allocator.wake_up(tags)
+        current_omni_platform.synchronize()
+        logger.info(f"[LLM Worker {self.rank}] Wake-up complete.")
+        return True
+
+    def handle_sleep_task(self, task: OmniSleepTask) -> OmniACK:
+        "Handle deterministic Sleep command from the main process"
+        try:
+            if isinstance(task, dict):
+                task = OmniSleepTask(**task)
+            logger.info(f"[LLM Worker {self.rank}] Handshake Received: Task {task.task_id}, Level {task.level}")
+            if task.level == 2:
+                if hasattr(self.model_runner, "graph_runners"):
+                    self.model_runner.graph_runners.clear()
+                    logger.info(f"[LLM Worker {self.rank}] LLM CUDA Graphs cleared.")
+            mem_before = current_omni_platform.get_current_memory_usage(self.device)
+            self.sleep(level=task.level)
+            mem_after = current_omni_platform.get_current_memory_usage(self.device)
+            rank_freed = max(0, mem_before - mem_after)
+            if torch.distributed.is_initialized():
+                t_freed = torch.tensor([float(rank_freed)], device=self.device)
+                torch.distributed.all_reduce(t_freed)
+                total_freed = int(t_freed.item())
+                torch.distributed.barrier()
+            else:
+                total_freed = rank_freed
+            if self.rank != 0:
+                return None
+            current_stage_id = getattr(self.vllm_config.model_config, "stage_id", 0)
+            ack = OmniACK(
+                task_id=task.task_id,
+                status="SUCCESS",
+                stage_id=current_stage_id,
+                rank=self.rank,
+                freed_bytes=total_freed,
+                metadata={
+                    "source": "omni_platform_audit",
+                    "total_freed_gib": f"{total_freed / 1024**3:.2f}",
+                    "rank_residual_gib": f"{mem_after / 1024**3:.2f}",
+                },
+            )
+            if hasattr(self, "result_mq") and self.result_mq:
+                self.result_mq.put(ack)
+            logger.info(f"[LLM Worker {self.rank}] ACK emitted for Task {task.task_id}")
+            return ack
+        except Exception as e:
+            logger.error(f"[LLM Worker {self.rank}] Sleep Task Failed: {e}", exc_info=True)
+            if torch.distributed.is_initialized():
+                try:
+                    torch.distributed.barrier()
+                except Exception:
+                    pass
+            return OmniACK(task_id=task.task_id, status="ERROR", error_msg=str(e))
+
+    def handle_wake_task(self, task: OmniWakeTask) -> OmniACK:
+        "Handle deterministic Wakeup command from the main process"
+        try:
+            if isinstance(task, dict):
+                task = OmniWakeTask(**task)
+            self.wake_up(tags=task.tags)
+            if torch.distributed.is_initialized():
+                torch.distributed.barrier()
+            gc.collect()
+            current_omni_platform.synchronize()
+            usage_now = current_omni_platform.get_current_memory_usage(self.device)
+            if self.rank != 0:
+                return None
+            current_stage_id = getattr(self.vllm_config.model_config, "stage_id", 0)
+            ack = OmniACK(
+                task_id=task.task_id,
+                status="SUCCESS",
+                stage_id=current_stage_id,
+                rank=self.rank,
+                metadata={"state": "WARM", "current_vram_gib": f"{usage_now / 1024**3:.2f}"},
+            )
+            if hasattr(self, "result_mq") and self.result_mq:
+                self.result_mq.put(ack)
+            logger.info(f"[LLM Worker {self.rank}] Wake-up ACK emitted.")
+            return ack
+        except Exception as e:
+            logger.error(f"[LLM Worker {self.rank}] Wake-up Failed: {e}", exc_info=True)
+            if torch.distributed.is_initialized():
+                try:
+                    torch.distributed.barrier()
+                except Exception:
+                    pass
+            tid = task.task_id if hasattr(task, "task_id") else "unknown"
+            return OmniACK(task_id=tid, status="ERROR", error_msg=str(e))

vllm_omni/worker/gpu_ar_worker.py

@@ -12,6 +12,7 @@ from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
 from vllm.v1.worker.utils import request_memory
 from vllm.v1.worker.workspace import init_workspace_manager
 
+from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask
 from vllm_omni.worker.base import OmniGPUWorkerBase
 from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
 from vllm_omni.worker.mixins import OmniWorkerMixin
@@ -104,3 +105,23 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
         if self.rank == 0:
             # If usage stat is enabled, collect relevant info.
             report_usage_stats(self.vllm_config)
+
+    def handle_sleep_task(self, task: OmniSleepTask | dict) -> OmniACK:
+        """
+        Explicitly handle sleep commands.
+        Calls the implementation in the base class OmniGPUWorkerBase.
+        """
+        logger.debug(f"[AR Worker {self.rank}] Resolving handle_sleep_task dispatch")
+        if isinstance(task, dict):
+            task = OmniSleepTask(**task)
+        return super().handle_sleep_task(task)
+
+    def handle_wake_task(self, task: OmniWakeTask | dict) -> OmniACK:
+        """
+        Explicitly handle wake-up commands.
+        Calls the implementation in the base class OmniGPUWorkerBase.
+        """
+        logger.debug(f"[AR Worker {self.rank}] Resolving handle_wake_task dispatch")
+        if isinstance(task, dict):
+            task = OmniWakeTask(**task)
+        return super().handle_wake_task(task)