update news

Signed-off-by: Roger Wang <hey@rogerw.io>

.buildkite/pipeline.yml

@@ -77,6 +77,24 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
+  - label: "Diffusion GPU Worker Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/diffusion/test_gpu_worker.py
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
   - label: "Omni Model Test"
     timeout_in_minutes: 15
     depends_on: image-build

README.md

@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 *Latest News* 🔥
 
+- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm).
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
 
 ---

docs/.nav.yml

@@ -49,7 +49,6 @@ nav:
     - design/architecture_overview.md
     - Feature Design:
       - design/feature/disaggregated_inference.md
-      - design/feature/multi_request_streaming.md
       - design/feature/ray_based_execution.md
     - Module Design:
       - design/module/ar_module.md

docs/configuration/README.md

@@ -2,7 +2,7 @@
 
 This section lists the most common options for running vLLM-Omni.
 
-For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.11.0/configuration/index.html)
+For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.12.0/configuration/index.html)
 
 Currently, the main options are maintained by stage configs for each model.
 

docs/contributing/profiling.md

@@ -1,3 +1,99 @@
 # Profiling vLLM-Omni
 
-Profiling is only intended for vLLM-Omni developers and maintainers to understand the proportion of time spent in different parts of the codebase. **vLLM-Omni end-users should never turn on profiling** as it will significantly slow down the inference.
+Performance Profiling Guidelines Profiling capabilities in vLLM-Omni are reserved for development and maintenance tasks aimed at temporal analysis of the codebase. Production use is **strongly discouraged**; enabling the profiler incurs a substantial overhead that negatively impacts inference latency.
+
+**Mechanism**: vLLM-Omni implements cross-stage profiling via the PyTorch Profiler. To accommodate the architecture—where stages operate as distinct engine instances in separate processes—the profiling interface supports both holistic capturing (all stages) and targeted capturing (specific stages).
+
+**1. Enabling the Profiler**
+
+Before running your script, you must set the ```VLLM_TORCH_PROFILER_DIR``` environment
+variable.
+
+```Bash
+export VLLM_TORCH_PROFILER_DIR=/path/to/save/traces
+```
+
+**Highly Recommended: Limit Profiling to a Single Iteration**  
+For most use cases (especially when profiling audio stages), you should limit the profiler to just **one iteration** to keep trace files small and readable.
+
+
+```bash
+export VLLM_PROFILER_MAX_ITERS=1
+```
+
+**2. Offline Inference**
+
+For offlinie processing using ```OmniLLM```, you can wrap your ```generate``` calls with ```start_profile``` and ```stop_profile()```.
+
+Basic Usage(All Stages)
+```Python
+from vllm_omni import OmniLLM
+
+omni_llm = OmniLLM.from_engine_args(engine_args)
+
+# Start profiling all active stages
+omni_llm.start_profile()
+
+outputs = omni_llm.generate(prompts, sampling_params)
+
+# Stop profiling and save traces
+omni_llm.stop_profile()
+```
+
+**Selective Stage Profiling**
+The profiler is default to function across all stages. But It is highly recommended to profile specific stages by passing the stages list, preventing from producing too large trace files:
+```python
+# Only profile Stage 1
+omni_llm.start_profile(stages=[1])
+```
+
+```python
+# Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
+omni_llm.start_profile(stages=[0, 2])
+```
+
+**Examples**:
+
+1. **Qwen-omni 2.5**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
+
+2. **Qwen-omni 3.0**:   [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
+
+**3. Online Inference(Async)**
+
+For online serving using AsyncOmni, the methods are asynchronous. This allows you to toggle profiling dynamically without restarting the server.
+
+```python
+from vllm_omni import AsyncOmni
+
+# Inside an async function:
+async_omni = AsyncOmni.from_engine_args(engine_args)
+
+await async_omni.start_profile()
+
+async for output in async_omni.generate(prompt, sampling_params, request_id):
+    # Process outputs...
+    pass
+
+await async_omni.stop_profile()
+```
+
+**4. Analyzing Omni Traces**
+
+After ``stop_profile()`` completes (and the file write wait time has passed), the directory specified in ```VLLM_TORCH_PROFILER_DIR``` will contain the trace files.
+
+```
+Output/
+│── ...rank-0.pt.trace.json.gz   # GPU 0 trace (TP=2 Example)
+│── ...rank-1.pt.trace.json.gz   # GPU 1 trace (TP=2 Example)
+│       # Load these into Perfetto to visualize synchronization
+│
+│── profiler_out_.txt            # Summary tables (CPU/CUDA time %)
+```
+
+**Viewing Tools:**
+     - [Perfetto](https://ui.perfetto.dev/): (Recommended): Best for handling large audio trace files.
+     - ```chrome://tracing```: Good for smaller text-only traces.
+
+
+**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM.  
+For more advanced configuration options (memory profiling, custom activities, etc.), see the official vLLM profiler documentation:  [vLLM Profiling Guide](https://docs.vllm.ai/en/latest/dev/profiling.html)

docs/design/feature/multi_request_streaming.md

@@ -1,47 +0,0 @@
-## Multi-Request Streaming (MRS) on a Single Machine
-
-### 1. Background & Scope
-- All processing runs on a single physical machine with multi-process, per-stage workers. No proxy or network transport involved.
-- Current alignment with vllm-omni: `OmniLLM` supports multiple stages (`OmniStage`). GPU runners already expose streamable steps (prefill/decoding/diffusion), but the entry layer still collects lists and lacks intra-stage streaming and window scheduling.
-- Goal: implement multi-stage, multi-request streaming (MRS) locally. Each stage outputs segments; downstream stages stitch and trigger compute based on configured windows. Shared memory and zero-copy strategies reduce data movement overhead.
-
-### 2. Key Constraints
-- Multi-process per stage: each stage is an independent process with a while loop; device visibility can be configured (`CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`).
-- Simple IPC (copy-based): use `multiprocessing.Queue`/Pipe for inter-process communication with CPU copies/serialization; do not rely on CUDA IPC/SHM zero-copy in this version.
-- Cross-stage pipeline: different stages can process different requests concurrently (e.g., stage A handles request 1 while stage B handles request 0).
-
-### 3. Architecture Overview
-- Processes & IPC queues
-  - Each "sub-stage" is an OS process (worker). The loop: take from input_queue → compute → put to output_queue.
-  - Inter-stage connection via IPC: copy-based `multiprocessing.Queue` passing dict payloads; use shared memory for large objects.
-  - Each link is SPSC (single-producer/single-consumer): the upstream is the orchestrator and the downstream is a single stage process; queues are unbounded (maxsize=0) on the orchestrator side.
-- Device visibility
-  - Each stage sets `CUDA_VISIBLE_DEVICES` or calls `torch.cuda.set_device` to bind to GPU sets.
-  - A stage may use multiple GPUs internally (TP/PP/DP) but presents as a single stage unit.
-- Simplified IPC: copy-based queues/pipes for data transfer; zero-copy is future work.
-- Pipeline progression: when a stage finishes a request, it enqueues outputs to the downstream stage; if downstream is idle, it starts immediately.
-- Scheduling
-  - A downstream stage triggers only after the upstream completes the request.
-  - Windowed segmentation/stitched triggering is not implemented; intra-stage streaming is not provided.
-
-### 4. IPC Implementation (simplified: copy-based)
-- Use `multiprocessing.Queue`/Pipe for inter-process communication (control + data).
-- Data is serialized/copied via CPU; no CUDA IPC/SHM zero-copy in this version.
-- Backpressure: queues are unbounded; pressure manifests as compute-rate differences. Optional SHM reduces large-object transfer cost; RX/decoding overhead is recorded for observability.
-
-### 5. Scheduling & Cancellation (simplified)
-- Pipeline: when a stage finishes a request, it enqueues to the next stage; that stage immediately pulls the next request from its input queue, enabling cross-stage concurrency.
-- Cancellation/timeout: explicit cancellation/timeouts are not provided; graceful shutdown uses a `None` sentinel sent to each stage input queue.
-
-#### Short sequence example (req0/req1, stage A→B)
-1) t0: stage A handles req0
-2) t1: req0 completes on A → enters B; A immediately starts req1
-3) t2: B handles req0 while A handles req1 (parallel across stages)
-
-### 6. Integration Points (by file)
-- `vllm_omni/entrypoints/omni.py` (Orchestrator)
-  - Class `Omni` orchestrates multi-process stages; constructs `OmniStage` instances in parallel and spawns per-stage workers.
-  - Spawns stage processes per config (set `CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`), creates control/data channels, builds simple full-trigger flow.
-  - Stats/logging are disabled by default; per-stage and orchestrator stats are only written when explicitly enabled.
-  - Manages process lifecycle: start/wait for readiness, graceful shutdown; forwards results between stages using copy-based IPC and optional SHM.
-  - Stage readiness: each stage emits `{"type": "stage_ready"}` after initialization; the orchestrator waits for all stages or times out and logs diagnostic suggestions.

docs/design/feature/ray_based_execution.md

@@ -1,14 +1,17 @@
 # Distributed utils
 
 This directory (vllm_omni/distributed/ray_utils) contains utilities for distributed execution in vllm-omni, supporting both **Ray** and **Multiprocessing** backends.
-
-## 1. Ray Utils
+## 1. Installation
+```bash
+pip install "ray[default]"
+```
+## 2. Ray Utils
 
 The `ray_utils` module provides helper functions for managing Ray clusters and actors, which is essential for:
 *   **Multi-node deployment**: Running pipeline stages across different physical machines.
 *   **Resource management**: Efficient GPU/CPU allocation.
 
-### 1.1 Basic Usage
+### 2.1 Basic Usage
 
 To use the Ray backend, specify `worker_backend="ray"` when initializing the engine.
 
@@ -21,7 +24,7 @@ vllm serve Qwen/Qwen2.5-Omni-7B \
   --ray-address auto
 ```
 
-### 1.2 Cluster Setup
+### 2.2 Cluster Setup
 
 **Step 1: Start Head Node**
 Run this on your primary machine:
@@ -38,7 +41,7 @@ ray start --address=<HEAD_NODE_IP>:6399
 > **Tip**: For a complete cluster setup script, refer to the vLLM example:
 > [run_cluster.sh](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/run_cluster.sh)
 
-### 1.3 Distributed Connector Support
+### 2.3 Distributed Connector Support
 
 When running on Ray, the system automatically adapts its communication strategy:
 
@@ -46,16 +49,11 @@ When running on Ray, the system automatically adapts its communication strategy:
 *   **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
 *   **SHM threshold default differs**: when `worker_backend="ray"`, the SharedMemoryConnector default threshold is set to `sys.maxsize`, which forces payloads to go inline (no SHM). Override `shm_threshold_bytes` in the connector config if you want SHM for Ray runs.
 
-### 1.4 Internal Helpers
+### 2.4 Internal Helpers
 
 *   **`initialize_ray_cluster`**: Connects to an existing Ray cluster or starts a local one.
 
-## 2. Troubleshooting
+## 3. Troubleshooting
 
 *   **Connection Issues**: Ensure the Ray head node is accessible and ports (default 6399 in this example) are open.
 *   **Version Mismatch**: Ensure all nodes run the same version of Ray and Python.
-
-### Installation
-```bash
-pip install "ray[default]"
-```

docs/design/index.md

@@ -9,7 +9,6 @@ This section contains design documents and architecture specifications for vLLM-
 ## Feature Design Documents
 
 - [Disaggregated Inference](feature/disaggregated_inference.md)
-- [Multi-Request Streaming](feature/multi_request_streaming.md)
 - [Ray-based Execution](feature/ray_based_execution.md)
 
 ## Module Design Documents

docs/design/module/dit_module.md

@@ -1,3 +1,7 @@
+---
+toc_depth: 4
+---
+
 # Diffusion Module Architecture Design
 
 The vLLM-Omni diffusion module (`vllm_omni/diffusion`) is a high-performance inference engine for diffusion models, designed with a modular architecture that separates concerns across multiple components. It provides efficient execution for non-autoregressive generation tasks such as image and video generation.
@@ -5,7 +9,7 @@ The vLLM-Omni diffusion module (`vllm_omni/diffusion`) is a high-performance inf
 This document describes the architecture design of the diffusion module, including the diffusion engine, scheduler, worker, diffusion pipeline, and acceleration components.
 
 <p align="center">
-   <img src="https://github.com/user-attachments/assets/cde1ebea-4006-4d30-8283-15a5721a09d2" alt="vLLM-Omni Diffusion Module Components" width="60%">
+   <img src="https://github.com/user-attachments/assets/cde1ebea-4006-4d30-8283-15a5721a09d2" alt="vLLM-Omni Diffusion Module Components" width="80%">
 </p>
 <p align="center">
   <em> Main Components of the Diffusion Module </em>
@@ -20,10 +24,10 @@ This document describes the architecture design of the diffusion module, includi
 - [Worker](#3-worker)
 - [Diffusion Pipeline](#4-diffusion-pipeline)
 - [Acceleration Components](#5-acceleration-components)
-   - [Attention Backends](#51-attention-backends)
-   - [Parallel Attention](#52-parallel-attention)
-   - [Cache Backends](#53-cache-backends)
-   - [Parallel Strategies](#54-parallel-strategies)
+    - [Attention Backends](#51-attention-backends)
+    - [Parallel Attention](#52-parallel-attention)
+    - [Cache Backends](#53-cache-backends)
+    - [Parallel Strategies](#54-parallel-strategies)
 - [Data Flow](#6-data-flow)
 
 ---
@@ -65,8 +69,11 @@ class DiffusionEngine:
 ```
 
 **Key Features**:
+
 - **Pre/Post Processing**: Registers model-specific pre-processing and post-processing functions via registry pattern
+
 - **Worker Management**: Launches and manages multiple worker processes (one per GPU)
+
 - **Process Isolation**: Uses multiprocessing for true parallelism
 
 #### 1.2 Worker Launch Process
@@ -86,8 +93,11 @@ def _launch_workers(self, broadcast_handle):
 ```
 
 **Design Decisions**:
+
 - **Spawn Method**: Ensures clean state for each worker (no shared memory issues)
+
 - **Pipe Communication**: Uses `mp.Pipe` for initialization handshake
+
 - **Device Selection**: Each worker is assigned a specific GPU (`cuda:{rank}`)
 
 #### 1.3 Request Processing Flow
@@ -106,8 +116,11 @@ def step(self, requests: list[OmniDiffusionRequest]):
 ```
 
 **Flow**:
+
 1. **Pre-processing**: Applies model-specific transformations
+
 2. **Scheduling**: Delegates to scheduler for distribution
+
 3. **Post-processing**: Converts raw outputs to final format (e.g., PIL images)
 
 ---
@@ -139,8 +152,11 @@ class Scheduler:
 ```
 
 **Communication Pattern**:
+
 - **Broadcast Queue**: One-to-many communication (scheduler → all workers)
+
 - **Result Queue**: One-to-one communication (rank 0 → scheduler)
+
 - **Shared Memory**: Uses `MessageQueue` (ZMQ-based) for efficient IPC
 
 #### 2.2 Request Distribution
@@ -156,8 +172,11 @@ def add_req(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
 ```
 
 **Design Features**:
+
 - **Broadcast Model**: All workers receive the same request (for tensor parallelism)
+
 - **Single Response**: Only rank 0 sends results back (avoids duplicate outputs)
+
 - **Synchronous**: Blocks until result is received (can be made async)
 
 #### 2.3 Singleton Pattern
@@ -176,8 +195,11 @@ scheduler = Scheduler()
 ```
 
 **Benefits**:
+
 - **Single Point of Control**: Ensures consistent state
+
 - **Easy Access**: Global `scheduler` instance accessible everywhere
+
 - **Resource Management**: Centralized queue management
 
 ---
@@ -212,12 +234,19 @@ class WorkerProc:
 ```
 
 **Initialization Steps**:
+
 1. **IPC Setup**: Creates ZMQ context and message queues
+
 2. **Distributed Environment Setup**: Initializes PyTorch distributed communication
-     - For CUDA GPUs: Uses NCCL (fast GPU communication)
-     - For NPU: Uses HCCL (Huawei Collective Communications Library)
-     - For other devices: Uses appropriate backend (GLOO, MCCL, etc.)
+
+    - For CUDA GPUs: Uses NCCL (fast GPU communication)
+
+    - For NPU: Uses HCCL (Huawei Collective Communications Library)
+
+    - For other devices: Uses appropriate backend (GLOO, MCCL, etc.)
+
 3. **Model Loading**: Loads diffusion pipeline on assigned GPU
+
 4. **Cache Setup**: Enables cache backend if configured.
 
 #### 3.2 GPU Worker
@@ -253,8 +282,11 @@ class GPUWorker:
 ```
 
 **Key Features**:
+
 - **Tensor Parallelism**: Supports multi-GPU tensor parallelism via PyTorch distributed
+
 - **Model Loading**: Uses `DiffusersPipelineLoader` for efficient weight loading
+
 - **Cache Integration**: Enables cache backends (TeaCache, cache-dit, etc.) transparently
 
 #### 3.3 Worker Busy Loop
@@ -283,9 +315,13 @@ def worker_busy_loop(self):
 ```
 
 **Execution Flow**:
+
 1. **Receive**: Dequeues unified messages from shared memory queue
+
 2. **Route**: Handles different message types (generation, RPC, shutdown)
+
 3. **Execute**: Runs forward pass through pipeline for generation requests
+
 4. **Respond**: Sends results back (rank 0 for generation, specified rank for RPC)
 
 #### 3.4 Model Execution
@@ -315,10 +351,14 @@ The model execution leverages multiple parallelism strategies that are transpare
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_sp_group, get_dp_group, get_cfg_group, get_pp_group
 )
+```
 
 **Optimizations**:
+
 - **Cache Refresh**: Clears cache state before each generation for clean state
+
 - **Context Management**: Forward context ensures parallel groups are available during execution
+
 - **Single Request**: Currently processes one request at a time (batching TODO)
 
 ---
@@ -327,66 +367,9 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 
 **Location**: `vllm_omni/diffusion/models/*/pipeline_*.py`
 
-### Architecture
-
 The pipeline is the **model-specific implementation** that orchestrates the diffusion process. Different models (QwenImage, Wan2.2, Z-Image) have their own pipeline implementations.
 
-### Key Components (QwenImage Example)
-
-#### 4.1 Pipeline Structure
-
-```python
-class QwenImagePipeline(nn.Module):
-    def __init__(self, od_config: OmniDiffusionConfig):
-        # Core components
-        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(...)
-        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(...)
-        self.vae = AutoencoderKLQwenImage.from_pretrained(...)
-        self.transformer = QwenImageTransformer2DModel(od_config=od_config)
-        self.tokenizer = Qwen2Tokenizer.from_pretrained(...)
-
-        # Cache backend (set by worker)
-        self._cache_backend = None
-```
-
-**Component Responsibilities**:
-- **Text Encoder**: Encodes prompts into embeddings
-- **Transformer**: Main denoising network (DiT-style)
-- **VAE**: Encodes/decodes between pixel and latent space
-- **Scheduler**: Manages timestep scheduling
-
-#### 4.2 Forward Pass
-
-```python
-def forward(self, req: OmniDiffusionRequest, ...) -> DiffusionOutput:
-    # 1. Encode prompts
-    prompt_embeds, prompt_embeds_mask = self.encode_prompt(prompt)
-
-    # 2. Prepare latents
-    latents = self.prepare_latents(batch_size, num_channels, height, width, ...)
-
-    # 3. Prepare timesteps
-    timesteps, num_inference_steps = self.prepare_timesteps(...)
-
-    # 4. Diffusion loop
-    latents = self.diffuse(
-        prompt_embeds, prompt_embeds_mask,
-        negative_prompt_embeds, negative_prompt_embeds_mask,
-        latents, img_shapes, txt_seq_lens,
-        timesteps, do_true_cfg, guidance, true_cfg_scale,
-    )
-
-    # 5. Decode latents
-    image = self.vae.decode(latents)
-
-    return DiffusionOutput(output=image)
-```
-
-The main steps of the forward pass are referred from `diffusers`.
-
-#### 4.3 Diffusion Loop
-
-The multi-step diffusion loop is usually the most time-consuming part during the overall inference process.
+Most pipeline implementation are referred from `diffusers`. The multi-step diffusion loop is usually the most time-consuming part during the overall inference process, which is defined by the `diffuse` function in the pipeline class. An example is as follows:
 
 ```python
 def diffuse(self, ...):
@@ -416,11 +399,14 @@ def diffuse(self, ...):
 ```
 
 **Key Features**:
+
 - **CFG Support**: Handles classifier-free guidance with separate forward passes
+
 - **Cache Branching**: Uses `cache_branch` parameter for cache-aware execution
+
 - **True CFG**: Implements advanced CFG with norm preservation
 
-To learn how to add a new diffusion pipeline, please view [Adding Diffusion Model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_diffusion_model)
+To learn more about the diffusion pipeline and how to add a new diffusion pipeline, please view [Adding Diffusion Model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_diffusion_model)
 
 ---
 
@@ -448,9 +434,13 @@ class Attention(nn.Module):
 ```
 
 **Available Backends**:
+
 - **FlashAttention**: Optimized CUDA kernel (FA2/FA3) - memory efficient via tiling
+
 - **SDPA**: PyTorch's scaled dot-product attention - default, cross-platform
+
 - **SageAttention**: Sparse attention implementation from SageAttention library
+
 - **AscendAttention**: NPU-optimized attention for Ascend hardware
 
 These backends provide the **kernel implementations** for attention computation. For attention-level sequence parallelism strategies (Ring Attention, Ulysses), see [Parallel Attention](#52-parallel-attention).
@@ -470,16 +460,25 @@ def get_attn_backend(head_size: int) -> type[AttentionBackend]:
 ```
 
 **Selection Priority**:
+
 1. **Environment Variable**: `DIFFUSION_ATTENTION_BACKEND` for manual override
-   - Valid values: `FLASH_ATTN`, `TORCH_SDPA`, `SAGE_ATTN`, `ASCEND`
-   - Example: `export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN`
+
+    - Valid values: `FLASH_ATTN`, `TORCH_SDPA`, `SAGE_ATTN`, `ASCEND`
+
+    - Example: `export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN`
+
 2. **Automatic Fallback**: Falls back to SDPA if selected backend unavailable
+
 3. **Hardware Detection**: Can select based on device type (NPU, CUDA, etc.)
 
 **Backend Availability**:
+
 - **SDPA**: Always available (PyTorch built-in)
+
 - **FlashAttention**: Requires `flash-attn` package installed
+
 - **SageAttention**: Requires `sage-attention` package (from THU-ML GitHub)
+
 - **AscendAttention**: Only available on Ascend NPU hardware
 
 #### Attention Backend Registry
@@ -563,8 +562,11 @@ Parallel attention strategies implement **Sequence Parallelism (SP) at the atten
 **Key Distinction**: Unlike AttentionBackend (which provides kernel implementations), ParallelAttentionStrategy provides communication patterns for multi-GPU attention parallelism. These strategies implement the `ParallelAttentionStrategy` interface and use AttentionBackend implementations internally.
 
 Both Ring Attention and Ulysses are forms of Sequence Parallelism (SP) that:
+
 - Split the sequence dimension across GPUs
+
 - Contribute to `sequence_parallel_size` (via `ring_degree` and `ulysses_degree`)
+
 - Work at the attention layer level (not model/pipeline level)
 
 #### Ulysses Sequence Parallelism (USP)
@@ -574,7 +576,9 @@ Both Ring Attention and Ulysses are forms of Sequence Parallelism (SP) that:
 USP is a sequence-parallel attention strategy that splits attention computation across multiple GPUs by distributing both the sequence dimension and attention heads. It uses **all-to-all communication** to efficiently parallelize attention for very long sequences. Specifically, it uses **all-to-all** collective operations to redistribute Q/K/V tensors before attention computation and gather results afterward.
 
 Ulysses splits attention computation in two dimensions:
+
 1. **Sequence Dimension**: Splits the sequence length across GPUs
+
 2. **Head Dimension**: Splits attention heads across GPUs
 
 **Configuration**: `ulysses_degree` contributes to `sequence_parallel_size`
@@ -601,9 +605,13 @@ class RingParallelAttention:
 ```
 
 **Integration**:
+
 - Ring Attention is activated when `ring_degree > 1` in parallel config
+
 - It's selected by `build_parallel_attention_strategy()` in the attention layer
+
 - The `Attention` layer routes to `_run_ring_attention()` when Ring is enabled
+
 - Works alongside attention backends: Ring handles communication, backends handle computation
 
 **Configuration**: `ring_degree` contributes to `sequence_parallel_size`
@@ -611,8 +619,11 @@ class RingParallelAttention:
 #### Relationship with AttentionBackend
 
 Parallel attention strategies (Ring, Ulysses) work **on top of** AttentionBackend implementations:
+
 - They use AttentionBackend for the actual attention computation (FlashAttention, SDPA, etc.)
+
 - They handle the multi-GPU communication/parallelization layer
+
 - They implement `ParallelAttentionStrategy` interface (not `AttentionBackend`)
 
 For general parallelism strategies (Data Parallelism, Tensor Parallelism, Pipeline Parallelism), see [Parallel Strategies](#54-parallel-strategies).
@@ -651,8 +662,11 @@ class CacheBackend(ABC):
 ```
 
 **Design Pattern**:
+
 - **Abstract Base Class**: Defines contract for all cache backends
+
 - **Pipeline-based**: Works with pipeline instances (not just transformers)
+
 - **State Management**: Provides refresh mechanism for clean state between generations
 
 #### Available Backends
@@ -686,9 +700,13 @@ class TeaCacheBackend(CacheBackend):
 ```
 
 **TeaCache Features**:
+
 - **Timestep-aware**: Caches based on timestep embedding similarity
+
 - **Adaptive**: Dynamically decides when to reuse cached computations
+
 - **CFG-aware**: Handles positive/negative branches separately
+
 - **Custom Hook System**: Uses a custom forward interception mechanism (via `HookRegistry`) that wraps the module's `forward` method, allowing transparent integration without modifying model code
 
 **2. Cache-DiT Backend**
@@ -710,9 +728,13 @@ class CacheDiTBackend(CacheBackend):
 ```
 
 **Cache-DiT Features**:
+
 - **DBCache**: Dynamic block caching with configurable compute blocks
+
 - **SCM**: Step Computation Masking for additional speedup
+
 - **TaylorSeer**: Advanced calibration for cache accuracy
+
 - **Dual-transformer Support**: Handles models like Wan2.2 with two transformers
 
 #### Cache Backend Selector
@@ -748,9 +770,13 @@ def get_cache_backend(
 ```
 
 **Usage Flow**:
+
 1. **Selection**: `get_cache_backend()` returns appropriate backend instance
+
 2. **Enable**: `backend.enable(pipeline)` called during worker initialization
+
 3. **Refresh**: `backend.refresh(pipeline, num_inference_steps)` called before each generation
+
 4. **Check**: `backend.is_enabled()` verifies cache is active
 
 ### 5.4 Parallel Strategies
@@ -762,23 +788,35 @@ def get_cache_backend(
 The system supports multiple orthogonal parallelism strategies:
 
 **Sequence Parallelism (SP)**
+
 - **Purpose**: Split sequence dimension across GPUs
+
 - **Attention-level SP**: Ring Attention and Ulysses (USP) implement SP at the attention layer level
-  - See [Parallel Attention](#52-parallel-attention) for details
-  - Configuration: `ulysses_degree` × `ring_degree` = `sequence_parallel_size`
+
+    - See [Parallel Attention](#52-parallel-attention) for details
+
+    - Configuration: `ulysses_degree` × `ring_degree` = `sequence_parallel_size`
+
 - **Use Case**: Very long sequences (e.g., high-resolution images)
 
 **Data Parallelism (DP)**
+
 - **Purpose**: Replicate model across GPUs, split batch
+
 - **Use Case**: Batch processing, throughput optimization
 
 **Tensor Parallelism (TP)** (Experimental)
+
 - **Purpose**: Split model weights across GPUs
+
 - **Implementation**: Uses vLLM's tensor parallel groups
+
 - **Use Case**: Large models that don't fit on single GPU
 
 **CFG Parallelism**  (under development)
+
 - **Purpose**: Parallelize Classifier-Free Guidance (positive/negative prompts)
+
 - **Infrastructure**: CFG parallel groups are initialized and available via `get_cfg_group()`
 
 #### Parallel Group Management

docs/getting_started/installation/gpu/cuda.inc.md

@@ -17,12 +17,12 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 # --8<-- [start:pre-built-wheels]
 
 #### Installation of vLLM
-Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1 and 0.12.0rc1. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.11.0 --torch-backend=auto
+uv pip install vllm==0.12.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -93,7 +93,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.11.0rc1 \
+    vllm/vllm-omni:v0.12.0rc1 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 

docs/getting_started/installation/npu/npu.inc.md

@@ -13,10 +13,10 @@ export DEVICE0=/dev/davinci0
 export DEVICE1=/dev/davinci1
 # Update the vllm-ascend image
 # Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1-a3
+export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \
@@ -34,12 +34,22 @@ docker run --rm \
     -p 8000:8000 \
     -it $IMAGE bash
 
+# Install the missing dependency of mooncake in the origin image.
+apt update
+apt install libjemalloc2
+echo "export LD_PRELOAD=/usr/lib/$(uname -m)-linux-gnu/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc
+source ~/.bashrc
+
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.11.0rc1 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.12.0rc1 https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 pip install -v -e .
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+# (Optional) Disable mooncake for stable capability
+mv /usr/local/Ascend/ascend-toolkit/latest/python/site-packages/mooncake \
+   /usr/local/Ascend/ascend-toolkit/latest/python/site-packages/mooncake.disabled
 ```
 
 The default workdir is `/workspace`, with vLLM, vLLM-Ascend and vLLM-Omni code placed in `/vllm-workspace` installed in development mode.

tests/diffusion/attention/test_sequence_parallel.py

@@ -12,13 +12,13 @@ from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.data import (
     DiffusionParallelConfig,
     OmniDiffusionConfig,
-    set_current_omni_diffusion_config,
 )
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
     init_distributed_environment,
     initialize_model_parallel,
 )
+from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.utils.platform_utils import detect_device_type
 
 device_type = detect_device_type()
@@ -385,7 +385,7 @@ def ulysses_attention_on_test_model(
     )
 
     # Set the config so Attention can access it
-    with set_current_omni_diffusion_config(od_config):
+    with set_forward_context(omni_diffusion_config=od_config):
         # Create model
         hidden_size = num_heads * head_size
 

tests/diffusion/test_gpu_worker.py

@@ -0,0 +1,261 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Unit tests for GPUWorker class.
+
+This module tests the GPUWorker implementation:
+- load_weights: Loading model weights
+- sleep: Putting worker into sleep mode (levels 1 and 2)
+- wake_up: Waking worker from sleep mode
+"""
+
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.worker.gpu_worker import GPUWorker
+
+
+@pytest.fixture
+def mock_od_config():
+    """Create a mock OmniDiffusionConfig."""
+    config = Mock()
+    config.num_gpus = 1
+    config.master_port = 12345
+    config.enable_sleep_mode = False
+    config.cache_backend = None
+    config.cache_config = None
+    config.model = "test-model"
+    return config
+
+
+@pytest.fixture
+def mock_gpu_worker(mock_od_config):
+    """Create a GPUWorker with mocked initialization."""
+    with patch.object(GPUWorker, "init_device_and_model"):
+        worker = GPUWorker(local_rank=0, rank=0, od_config=mock_od_config)
+        # Mock the pipeline
+        worker.pipeline = Mock()
+        worker.cache_backend = None
+        return worker
+
+
+class TestGPUWorkerLoadWeights:
+    """Test GPUWorker.load_weights method."""
+
+    def test_load_weights_calls_pipeline(self, mock_gpu_worker):
+        """Test that load_weights delegates to pipeline.load_weights."""
+        # Setup mock weights
+        mock_weights = [
+            ("layer1.weight", torch.randn(10, 10)),
+            ("layer2.weight", torch.randn(20, 20)),
+        ]
+        expected_loaded = {"layer1.weight", "layer2.weight"}
+
+        # Configure pipeline mock
+        mock_gpu_worker.pipeline.load_weights = Mock(return_value=expected_loaded)
+
+        # Call load_weights
+        result = mock_gpu_worker.load_weights(mock_weights)
+
+        # Verify pipeline.load_weights was called with the weights
+        mock_gpu_worker.pipeline.load_weights.assert_called_once_with(mock_weights)
+        assert result == expected_loaded
+
+    def test_load_weights_empty_iterable(self, mock_gpu_worker):
+        """Test load_weights with empty weights iterable."""
+        mock_gpu_worker.pipeline.load_weights = Mock(return_value=set())
+
+        result = mock_gpu_worker.load_weights([])
+
+        mock_gpu_worker.pipeline.load_weights.assert_called_once_with([])
+        assert result == set()
+
+
+class TestGPUWorkerSleep:
+    """Test GPUWorker.sleep method."""
+
+    @patch("vllm_omni.diffusion.worker.gpu_worker.torch.cuda.mem_get_info")
+    @patch("vllm.device_allocator.cumem.CuMemAllocator")
+    def test_sleep_level_1(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
+        """Test sleep mode level 1 (offload weights only)."""
+        # Setup memory info mocks
+        # Before sleep: 1GB free, 8GB total
+        # After sleep: 3GB free, 8GB total (freed 2GB)
+        mock_mem_info.side_effect = [
+            (1 * 1024**3, 8 * 1024**3),  # Before sleep
+            (3 * 1024**3, 8 * 1024**3),  # After sleep
+        ]
+
+        # Setup allocator mock
+        mock_allocator = Mock()
+        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
+        mock_allocator.sleep = Mock()
+
+        # 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
+        # Verify buffers were NOT saved (level 1 doesn't save buffers)
+        assert len(mock_gpu_worker._sleep_saved_buffers) == 0
+
+    @patch("vllm_omni.diffusion.worker.gpu_worker.torch.cuda.mem_get_info")
+    @patch("vllm.device_allocator.cumem.CuMemAllocator")
+    def test_sleep_level_2(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
+        """Test sleep mode level 2 (offload all, save buffers)."""
+        # Setup memory info mocks
+        mock_mem_info.side_effect = [
+            (1 * 1024**3, 8 * 1024**3),  # Before sleep
+            (5 * 1024**3, 8 * 1024**3),  # After sleep (freed 4GB)
+        ]
+
+        # Setup allocator mock
+        mock_allocator = Mock()
+        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
+        mock_allocator.sleep = Mock()
+
+        # Mock pipeline buffers
+        mock_buffer1 = torch.randn(10, 10)
+        mock_buffer2 = torch.randn(20, 20)
+        mock_gpu_worker.pipeline.named_buffers = Mock(
+            return_value=[
+                ("buffer1", mock_buffer1),
+                ("buffer2", mock_buffer2),
+            ]
+        )
+
+        # Call sleep with level 2
+        result = mock_gpu_worker.sleep(level=2)
+
+        # Verify sleep was called with empty tags (offload all)
+        mock_allocator.sleep.assert_called_once_with(offload_tags=tuple())
+        assert result is True
+
+        # Verify buffers were saved
+        assert len(mock_gpu_worker._sleep_saved_buffers) == 2
+        assert "buffer1" in mock_gpu_worker._sleep_saved_buffers
+        assert "buffer2" in mock_gpu_worker._sleep_saved_buffers
+
+    @patch("vllm_omni.diffusion.worker.gpu_worker.torch.cuda.mem_get_info")
+    @patch("vllm.device_allocator.cumem.CuMemAllocator")
+    def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
+        """Test that sleep validates memory was actually freed."""
+        # Simulate memory increase (should trigger assertion error)
+        mock_mem_info.side_effect = [
+            (3 * 1024**3, 8 * 1024**3),  # Before sleep: 3GB free
+            (1 * 1024**3, 8 * 1024**3),  # After sleep: 1GB free (negative freed!)
+        ]
+
+        mock_allocator = Mock()
+        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
+        mock_allocator.sleep = Mock()
+
+        # This should raise an assertion error
+        with pytest.raises(AssertionError, match="Memory usage increased after sleeping"):
+            mock_gpu_worker.sleep(level=1)
+
+
+class TestGPUWorkerWakeUp:
+    """Test GPUWorker.wake_up method."""
+
+    @patch("vllm.device_allocator.cumem.CuMemAllocator")
+    def test_wake_up_without_buffers(self, mock_allocator_class, mock_gpu_worker):
+        """Test wake_up without saved buffers (level 1 sleep)."""
+        # Setup allocator mock
+        mock_allocator = Mock()
+        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
+        mock_allocator.wake_up = Mock()
+
+        # Ensure no saved buffers
+        mock_gpu_worker._sleep_saved_buffers = {}
+
+        # Call wake_up
+        result = mock_gpu_worker.wake_up(tags=["weights"])
+
+        # Verify allocator.wake_up was called
+        mock_allocator.wake_up.assert_called_once_with(["weights"])
+        assert result is True
+
+    @patch("vllm.device_allocator.cumem.CuMemAllocator")
+    def test_wake_up_with_buffers(self, mock_allocator_class, mock_gpu_worker):
+        """Test wake_up with saved buffers (level 2 sleep)."""
+        # Setup allocator mock
+        mock_allocator = Mock()
+        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
+        mock_allocator.wake_up = Mock()
+
+        # Create saved buffers
+        saved_buffer1 = torch.randn(10, 10)
+        saved_buffer2 = torch.randn(20, 20)
+        mock_gpu_worker._sleep_saved_buffers = {
+            "buffer1": saved_buffer1,
+            "buffer2": saved_buffer2,
+        }
+
+        # Mock pipeline buffers (these will be restored)
+        mock_buffer1 = Mock()
+        mock_buffer1.data = Mock()
+        mock_buffer2 = Mock()
+        mock_buffer2.data = Mock()
+
+        mock_gpu_worker.pipeline.named_buffers = Mock(
+            return_value=[
+                ("buffer1", mock_buffer1),
+                ("buffer2", mock_buffer2),
+            ]
+        )
+
+        # Call wake_up
+        result = mock_gpu_worker.wake_up(tags=None)
+
+        # Verify allocator.wake_up was called
+        mock_allocator.wake_up.assert_called_once_with(None)
+
+        # Verify buffers were restored
+        mock_buffer1.data.copy_.assert_called_once()
+        mock_buffer2.data.copy_.assert_called_once()
+
+        # Verify saved buffers were cleared
+        assert len(mock_gpu_worker._sleep_saved_buffers) == 0
+        assert result is True
+
+    @patch("vllm.device_allocator.cumem.CuMemAllocator")
+    def test_wake_up_partial_buffer_restore(self, mock_allocator_class, mock_gpu_worker):
+        """Test wake_up only restores buffers that were saved."""
+        # Setup allocator mock
+        mock_allocator = Mock()
+        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
+        mock_allocator.wake_up = Mock()
+
+        # Only save buffer1, not buffer2
+        saved_buffer1 = torch.randn(10, 10)
+        mock_gpu_worker._sleep_saved_buffers = {
+            "buffer1": saved_buffer1,
+        }
+
+        # Mock pipeline has both buffers
+        mock_buffer1 = Mock()
+        mock_buffer1.data = Mock()
+        mock_buffer2 = Mock()
+        mock_buffer2.data = Mock()
+
+        mock_gpu_worker.pipeline.named_buffers = Mock(
+            return_value=[
+                ("buffer1", mock_buffer1),
+                ("buffer2", mock_buffer2),
+            ]
+        )
+
+        # Call wake_up
+        result = mock_gpu_worker.wake_up()
+
+        # Verify only buffer1 was restored
+        mock_buffer1.data.copy_.assert_called_once()
+        # buffer2 should NOT be restored since it wasn't saved
+        mock_buffer2.data.copy_.assert_not_called()
+
+        assert result is True

vllm_omni/diffusion/attention/layer.py

@@ -15,8 +15,8 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy
 from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
 from vllm_omni.diffusion.attention.selector import get_attn_backend
-from vllm_omni.diffusion.data import get_current_omni_diffusion_config
 from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
+from vllm_omni.diffusion.forward_context import get_forward_context
 
 logger = init_logger(__name__)
 
@@ -58,7 +58,7 @@ class Attention(nn.Module):
         self.ring_runner = None
 
         try:
-            config = get_current_omni_diffusion_config()
+            config = get_forward_context().omni_diffusion_config
             self.backend_pref = config.attention_backend
             if config.parallel_config.ring_degree > 1:
                 self.use_ring = True

vllm_omni/diffusion/attention/parallel/factory.py

@@ -6,8 +6,8 @@ from __future__ import annotations
 from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention, ParallelAttentionStrategy
 from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
 from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
-from vllm_omni.diffusion.data import get_current_omni_diffusion_config
 from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size, get_sp_group
+from vllm_omni.diffusion.forward_context import get_forward_context
 
 
 def build_parallel_attention_strategy(
@@ -24,7 +24,7 @@ def build_parallel_attention_strategy(
       and initialized process groups.
     """
     try:
-        cfg = get_current_omni_diffusion_config()
+        cfg = get_forward_context().omni_diffusion_config
         p = cfg.parallel_config
     except Exception:
         return NoParallelAttention()

vllm_omni/diffusion/attention/parallel/ring.py

@@ -16,10 +16,10 @@ from vllm_omni.diffusion.attention.parallel.base import (
     ParallelAttentionContext,
     # ParallelAttentionStrategy, # Not used in type hint below currently
 )
+from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGroupCoordinator
 
 # from vllm_omni.diffusion.attention.backends.ring_selector import AttnType # Already imported above
-from vllm_omni.diffusion.data import get_current_omni_diffusion_config
-from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGroupCoordinator
+from vllm_omni.diffusion.forward_context import get_forward_context
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
@@ -110,7 +110,7 @@ class RingParallelAttention:
         backend_pref = self.attn_backend_pref
         if backend_pref is None:
             try:
-                config = get_current_omni_diffusion_config()
+                config = get_forward_context().omni_diffusion_config
                 # config might not have attention_backend attribute if not updated
                 backend_pref = getattr(config, "attention_backend", None)
             except Exception:

vllm_omni/diffusion/data.py

@@ -5,9 +5,7 @@ import enum
 import os
 import random
 from collections.abc import Callable
-from contextlib import contextmanager
 from dataclasses import dataclass, field, fields
-from functools import lru_cache
 from typing import Any
 
 import torch
@@ -307,6 +305,9 @@ class OmniDiffusionConfig:
     # Compilation
     enable_torch_compile: bool = False
 
+    # Enable sleep mode
+    enable_sleep_mode: bool = False
+
     disable_autocast: bool = False
 
     # VSA parameters
@@ -455,56 +456,6 @@ class OmniDiffusionConfig:
         return cls(**kwargs)
 
 
-_current_omni_diffusion_config: OmniDiffusionConfig | None = None
-_current_prefix: str | None = None
-
-
-@contextmanager
-def set_current_omni_diffusion_config(
-    omni_diffusion_config: OmniDiffusionConfig, check_compile=False, prefix: str | None = None
-):
-    """
-    Temporarily set the current vLLM-Omni config.
-    Used during model initialization.
-    We save the current vLLM-Omni config in a global variable,
-    so that all modules can access it, e.g. custom ops
-    can access the vLLM-Omni config to determine how to dispatch.
-    """
-    global _current_omni_diffusion_config, _current_prefix
-    old_omni_diffusion_config = _current_omni_diffusion_config
-    old_prefix = _current_prefix
-    # from vllm.compilation.counter import compilation_counter
-
-    # num_models_seen = compilation_counter.num_models_seen
-    try:
-        _current_omni_diffusion_config = omni_diffusion_config
-        _current_prefix = prefix
-        yield
-    except Exception:
-        raise
-    else:
-        if check_compile:
-            raise RuntimeError("Compilation is not yet supported for OmniDiffusion")
-    finally:
-        _current_omni_diffusion_config = old_omni_diffusion_config
-        _current_prefix = old_prefix
-        # Clear the compilation config cache when context changes
-        get_cached_compilation_config.cache_clear()
-
-
-@lru_cache(maxsize=1)
-def get_cached_compilation_config():
-    """Cache config to avoid repeated calls to get_current_omni_diffusion_config()"""
-    return get_current_omni_diffusion_config().compilation_config
-
-
-def get_current_omni_diffusion_config() -> OmniDiffusionConfig:
-    if _current_omni_diffusion_config is None:
-        logger.warning_once("Current OmniDiffusionConfig is not set.")
-        return OmniDiffusionConfig()
-    return _current_omni_diffusion_config
-
-
 @dataclass
 class DiffusionOutput:
     """

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -24,7 +24,7 @@ from vllm_omni.diffusion.attention.backends.abstract import (
 )
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.cache.base import CachedTransformer
-from vllm_omni.diffusion.data import OmniDiffusionConfig, get_current_omni_diffusion_config
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_sequence_parallel_rank,
     get_sequence_parallel_world_size,
@@ -362,7 +362,7 @@ class QwenImageCrossAttention(nn.Module):
         self.rope = RotaryEmbedding(is_neox_style=False)
 
         try:
-            config = get_current_omni_diffusion_config()
+            config = get_forward_context().omni_diffusion_config
             self.parallel_config = config.parallel_config
         except Exception:
             self.parallel_config = None

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -181,16 +181,30 @@ class Wan22Pipeline(nn.Module):
 
         # Read model_index.json to detect expand_timesteps mode (for TI2V-5B)
         self.expand_timesteps = False
+        self.has_transformer_2 = False
         if local_files_only:
             model_index_path = os.path.join(model, "model_index.json")
             if os.path.exists(model_index_path):
                 with open(model_index_path) as f:
                     model_index = json.load(f)
                     self.expand_timesteps = model_index.get("expand_timesteps", False)
+            # Check if this is a two-stage model (MoE with transformer_2)
+            transformer_2_path = os.path.join(model, "transformer_2")
+            self.has_transformer_2 = os.path.exists(transformer_2_path)
+        else:
+            # For remote models, download and read model_index.json
+            try:
+                from huggingface_hub import hf_hub_download
 
-        # Check if this is a two-stage model (MoE with transformer_2)
-        transformer_2_path = os.path.join(model, "transformer_2") if local_files_only else None
-        self.has_transformer_2 = transformer_2_path is not None and os.path.exists(transformer_2_path)
+                model_index_path = hf_hub_download(repo_id=model, filename="model_index.json")
+                with open(model_index_path) as f:
+                    model_index = json.load(f)
+                    self.expand_timesteps = model_index.get("expand_timesteps", False)
+                    # Check transformer_2 from model_index
+                    transformer_2_info = model_index.get("transformer_2", [None, None])
+                    self.has_transformer_2 = transformer_2_info[0] is not None
+            except Exception:
+                pass
 
         # Set up weights sources for transformer(s)
         self.weights_sources = [
@@ -222,17 +236,13 @@ class Wan22Pipeline(nn.Module):
         ).to(self.device)
 
         # Initialize transformers with correct config (weights loaded via load_weights)
-        if local_files_only:
-            transformer_config = load_transformer_config(model, "transformer")
-            self.transformer = create_transformer_from_config(transformer_config)
-            if self.has_transformer_2:
-                transformer_2_config = load_transformer_config(model, "transformer_2")
-                self.transformer_2 = create_transformer_from_config(transformer_2_config)
-            else:
-                self.transformer_2 = None
+        transformer_config = load_transformer_config(model, "transformer", local_files_only)
+        self.transformer = create_transformer_from_config(transformer_config)
+        if self.has_transformer_2:
+            transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only)
+            self.transformer_2 = create_transformer_from_config(transformer_2_config)
         else:
-            self.transformer = WanTransformer3DModel()
-            self.transformer_2 = WanTransformer3DModel() if self.has_transformer_2 else None
+            self.transformer_2 = None
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py

@@ -31,8 +31,11 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import retrieve_latents
-from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
+    create_transformer_from_config,
+    load_transformer_config,
+    retrieve_latents,
+)
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 
@@ -149,7 +152,9 @@ class Wan22TI2VPipeline(nn.Module):
         ).to(self.device)
 
         # Single transformer (TI2V uses dense 5B model, not MoE)
-        self.transformer = WanTransformer3DModel()
+        # Load config from model to get correct dimensions
+        transformer_config = load_transformer_config(model, "transformer", local_files_only)
+        self.transformer = create_transformer_from_config(transformer_config)
 
         # Scheduler
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(

vllm_omni/diffusion/worker/gpu_worker.py

@@ -3,10 +3,12 @@
 import multiprocessing as mp
 import os
 import time
+from collections.abc import Iterable
+from contextlib import AbstractContextManager, nullcontext
 
 import torch
 import zmq
-from vllm.config import LoadConfig, VllmConfig, set_current_vllm_config
+from vllm.config import LoadConfig, VllmConfig
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.logger import init_logger
 from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes
@@ -15,7 +17,6 @@ from vllm_omni.diffusion.cache.selector import get_cache_backend
 from vllm_omni.diffusion.data import (
     DiffusionOutput,
     OmniDiffusionConfig,
-    set_current_omni_diffusion_config,
 )
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
@@ -45,7 +46,7 @@ class GPUWorker:
         self.od_config = od_config
         self.pipeline = None
         self.device = None
-
+        self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
         self.init_device_and_model()
 
     def init_device_and_model(self) -> None:
@@ -67,10 +68,7 @@ class GPUWorker:
         vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size
         vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size
         self.vllm_config = vllm_config
-        with (
-            set_current_omni_diffusion_config(self.od_config),
-            set_current_vllm_config(vllm_config),
-        ):
+        with set_forward_context(vllm_config=vllm_config, omni_diffusion_config=self.od_config):
             init_distributed_environment(world_size=world_size, rank=rank)
             logger.info(f"Worker {self.rank}: Initialized device and distributed environment.")
             parallel_config = self.od_config.parallel_config
@@ -87,11 +85,12 @@ class GPUWorker:
             load_config = LoadConfig()
             model_loader = DiffusersPipelineLoader(load_config)
             time_before_load = time.perf_counter()
-            with DeviceMemoryProfiler() as m:
-                self.pipeline = model_loader.load_model(
-                    od_config=self.od_config,
-                    load_device=str(self.device),
-                )
+            with self._maybe_get_memory_pool_context(tag="weights"):
+                with DeviceMemoryProfiler() as m:
+                    self.pipeline = model_loader.load_model(
+                        od_config=self.od_config,
+                        load_device=str(self.device),
+                    )
             time_after_load = time.perf_counter()
 
         logger.info(
@@ -140,6 +139,58 @@ class GPUWorker:
             output = self.pipeline.forward(req)
         return output
 
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        return self.pipeline.load_weights(weights)
+
+    def sleep(self, level: int = 1) -> bool:
+        from vllm.device_allocator.cumem import CuMemAllocator
+
+        free_bytes_before_sleep = torch.cuda.mem_get_info()[0]
+
+        # Save the buffers before level 2 sleep
+        if level == 2:
+            model = self.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())
+        free_bytes_after_sleep, total = torch.cuda.mem_get_info()
+        freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
+        used_bytes = total - free_bytes_after_sleep
+        assert freed_bytes >= 0, "Memory usage increased after sleeping."
+        logger.info(
+            "Sleep mode freed %.2f GiB memory, %.2f GiB memory is still in use.",
+            freed_bytes / GiB_bytes,
+            used_bytes / GiB_bytes,
+        )
+        return True
+
+    def wake_up(self, tags: list[str] | None = None) -> bool:
+        from vllm.device_allocator.cumem import CuMemAllocator
+
+        allocator = CuMemAllocator.get_instance()
+        allocator.wake_up(tags)
+
+        # Restore the buffers after level 2 sleep
+        if len(self._sleep_saved_buffers):
+            model = self.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 = {}
+        return True
+
+    def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
+        if self.od_config.enable_sleep_mode:
+            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."
+            return allocator.use_memory_pool(tag=tag)
+        else:
+            return nullcontext()
+
     def shutdown(self) -> None:
         destroy_distributed_env()
 

vllm_omni/diffusion/worker/npu/npu_worker.py

@@ -5,19 +5,19 @@ import os
 import time
 
 import torch
-from vllm.config import LoadConfig, VllmConfig, set_current_vllm_config
+from vllm.config import LoadConfig, VllmConfig
 from vllm.logger import init_logger
 from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes
 
 from vllm_omni.diffusion.cache.selector import get_cache_backend
 from vllm_omni.diffusion.data import (
     OmniDiffusionConfig,
-    set_current_omni_diffusion_config,
 )
 from vllm_omni.diffusion.distributed.parallel_state import (
     init_distributed_environment,
     initialize_model_parallel,
 )
+from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.worker.gpu_worker import GPUWorker, WorkerProc
 
@@ -49,10 +49,7 @@ class NPUWorker(GPUWorker):
         vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size
         vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size
         self.vllm_config = vllm_config
-        with (
-            set_current_omni_diffusion_config(self.od_config),
-            set_current_vllm_config(vllm_config),
-        ):
+        with set_forward_context(vllm_config=vllm_config, omni_diffusion_config=self.od_config):
             init_distributed_environment(world_size=world_size, rank=rank)
             logger.info(f"Worker {self.rank}: Initialized device and distributed environment.")
             parallel_config = self.od_config.parallel_config

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py

@@ -101,6 +101,12 @@ class Qwen2_5OmniForConditionalGeneration(
             self.talker.init_multi_modal(thinker_config)
             self.model = self.talker
             self.token2wav = None
+            # set suppress start id according to token2wav
+            t2w_token_end_id = getattr(
+                getattr(getattr(config, "token2wav_config", None), "dit_config", None), "num_embeds", None
+            )
+            if t2w_token_end_id:
+                self.model.set_suppress_start_id(t2w_token_end_id + 1)
 
         elif self.model_stage == "code2wav":
             self.thinker = None

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py

@@ -84,6 +84,9 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
         )
         self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
 
+        # suppress start id
+        self.suppress_start_id = None
+
     def init_multi_modal(self, thinker_config):
         self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
         self.visual = Qwen2_5_VisionTransformer(
@@ -150,6 +153,10 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
         return hidden_states
 
     def bad_word_processor(self, logits: torch.Tensor) -> torch.Tensor:
+        # suppress token IDs unsupported by token2wav
+        if self.suppress_start_id and self.suppress_start_id < logits.size(-1):
+            logits[..., self.suppress_start_id : logits.size(-1)] = -1e9
+
         if hasattr(self.config, "tts_codec_start_token_id"):
             bos_id = int(getattr(self.config, "tts_codec_start_token_id"))
             logits[..., bos_id] = -1e9
@@ -234,3 +241,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
                 audio_embeddings = self._process_audio_input(multimodal_input)
                 multimodal_embeddings += audio_embeddings
         return multimodal_embeddings
+
+    def set_suppress_start_id(self, start_id: int):
+        self.suppress_start_id = start_id
+        self.logger.debug(f"Set suppress start id to {self.suppress_start_id}")