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—without stopping the server or unloading the Docker container.

This feature is inherited from vLLM’s Sleep Mode 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.

1. Feature Documentation¶

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.

Sleep Levels¶

We support two levels of hibernation to balance recovery speed and memory efficiency:

Supported Platforms¶

Omni Sleep Mode is optimized for high-performance computing backends:

• 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.

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.

Quiescing a busy diffusion stage: pause_generation(mode="keep")¶

sleep() on a diffusion stage is a worker-level RPC: it offloads memory but does not stop the stage's scheduler, so a stage that still has queued requests must be made quiet first. Aborting discards the batch that is running. pause_generation(mode="keep") keeps it instead:

• The batch that is already running finishes on the weights it started with and its outputs are delivered normally.
• Requests that have not started stay queued and run after resume_generation(), on whatever weights are loaded by then. A request never continues denoising across a weight change.
• The call returns only after every replica has acknowledged: the batch that was running has finished executing, its asynchronous output copies are done, and the device is synchronized; its outputs are then delivered on the normal path. On a non-preemptible backend that acknowledgement arrives at the next batch boundary, so expect it to take as long as the remaining batch.
• Control RPCs (list_loras, sleep, wake_up, resume_generation) keep working while paused.

Recommended order for a weight update on a diffusion stage:

await engine.pause_generation(mode="keep", clear_cache=False)
await engine.sleep(level=1)          # optional: free memory for the trainer
await engine.wake_up()               # admission stays closed until resume
# install new weights here
await engine.resume_generation()

Support matrix and limits:

Behavior notes:

• Callers that already passed mode="keep" for a diffusion stage used to get an immediate return that only closed frontend admission. They now also stop the backend and wait for the ACK.
• A timeout or error from pause_generation means the pause did not complete; the scheduler stays closed and the caller must retry the pause or call resume_generation() explicitly. Do not sleep or mutate weights after a failed pause.
• resume_generation(stage_ids=...) that leaves another paused stage closed (an AR stage, or a keep-paused diffusion stage) keeps frontend admission closed as well, so new requests cannot queue on a stage that will not schedule them. Admission reopens once every stage paused by pause_generation has been resumed.
• clear_cache has no diffusion-specific effect.
• Without a pause call nothing changes: no barrier runs and batching, abort and output ordering are as before.

2. Usage Examples¶

Python API Example¶

You can programmatically control the lifecycle of stages using the AsyncOmni engine.

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¶

OmniACK Dataclass Fields¶

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.