From 89f8819525589141fd825ce4f0d1e1be9cf3660b Mon Sep 17 00:00:00 2001
From: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
Date: Tue, 19 May 2026 09:42:17 +0800
Subject: [PATCH] [Feature] Add support for Pipeline Parallel and integrate it
 into Wan 2.2 (#2322)

Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
---
 docs/design/feature/pipeline_parallel.md      | 356 ++++++++
 .../diffusion/frame_interpolation.md          |   1 -
 .../diffusion/parallelism/overview.md         |  17 +-
 .../parallelism/pipeline_parallel.md          | 205 +++++
 docs/user_guide/diffusion_features.md         |  95 ++-
 .../image_to_video/image_to_video.py          |  76 +-
 .../text_to_video/text_to_video.py            |  60 +-
 .../distributed/test_pipeline_parallel.py     | 802 ++++++++++++++++++
 .../models/wan2_2/test_wan22_ti2v_pipeline.py |  97 ---
 .../diffusion/cache/cache_dit_backend.py      |   1 -
 .../diffusion/distributed/cfg_parallel.py     |   8 +-
 .../distributed/group_coordinator.py          |  72 ++
 .../diffusion/distributed/parallel_state.py   |   5 -
 .../distributed/pipeline_parallel.py          | 291 +++++++
 vllm_omni/diffusion/models/wan2_2/__init__.py |   8 -
 .../models/wan2_2/pipeline_wan2_2.py          |  19 +-
 .../models/wan2_2/pipeline_wan2_2_i2v.py      |  23 +-
 .../models/wan2_2/pipeline_wan2_2_ti2v.py     | 731 ----------------
 .../models/wan2_2/wan2_2_transformer.py       | 122 ++-
 vllm_omni/engine/async_omni_engine.py         |   1 +
 20 files changed, 1981 insertions(+), 1009 deletions(-)
 create mode 100644 docs/design/feature/pipeline_parallel.md
 create mode 100644 docs/user_guide/diffusion/parallelism/pipeline_parallel.md
 create mode 100644 tests/diffusion/distributed/test_pipeline_parallel.py
 delete mode 100644 tests/diffusion/models/wan2_2/test_wan22_ti2v_pipeline.py
 create mode 100644 vllm_omni/diffusion/distributed/pipeline_parallel.py
 delete mode 100644 vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py

diff --git a/docs/design/feature/pipeline_parallel.md b/docs/design/feature/pipeline_parallel.md
new file mode 100644
index 000000000..3ed030288
--- /dev/null
+++ b/docs/design/feature/pipeline_parallel.md
@@ -0,0 +1,356 @@
+# Pipeline Parallel
+
+This section describes how to add Pipeline Parallelism (PP) to a diffusion pipeline. We use the Wan2.2 text-to-video and
+image-to-video pipelines as the reference implementations.
+
+## Implementation Checklist
+
+Adding Pipeline Parallel support requires:
+
+1. ✅ **Inherit the mixin** - add `PipelineParallelMixin` before `CFGParallelMixin` in the pipeline class
+2. ✅ **Make stage forward resumable** - support `intermediate_tensors` in the transformer forward path
+3. ✅ **Return the right object type** - `IntermediateTensors` on non-last PP ranks, final model output on the last rank
+4. ✅ **Use the standard denoising loop** - call `predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()`
+5. ✅ **Test parity** - compare PP results against the single-GPU baseline
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Troubleshooting](#troubleshooting)
+- [Reference Implementations](#reference-implementations)
+
+---
+
+## Overview
+
+### What is Pipeline Parallelism?
+
+Pipeline Parallelism splits the denoising transformer into multiple sequential stages and places each stage on a
+different rank. Instead of every rank holding the full DiT, each PP rank owns only a slice of the layers.
+
+For each denoising step:
+
+1. Rank 0 starts the forward pass using the current latents.
+2. Each intermediate rank receives hidden states from the previous rank, runs its local layer slice, and forwards the
+   intermediate tensors downstream.
+3. The last PP rank produces the final noise prediction.
+4. The last PP rank applies the scheduler step and sends the updated latents back to rank 0 for the next timestep.
+
+This reduces per-rank model memory and enables larger diffusion transformers to run across multiple GPUs. It can also be
+combined with CFG-Parallel, where each PP pipeline carries one CFG branch.
+
+### Architecture
+
+vLLM-Omni provides `PipelineParallelMixin` to encapsulate the PP communication pattern for diffusion pipelines.
+
+| Method                            | Purpose                       | Automatic Behavior                                                              |
+|-----------------------------------|-------------------------------|---------------------------------------------------------------------------------|
+| `diffuse()`                       | Denoising loop boundary       | Wrapped by `PipelineParallelMixin` to flush pending async PP sends on exit      |
+| `predict_noise_maybe_with_cfg()`  | Predict noise with PP support | Runs partial forwards on non-last PP ranks, combines with CFG logic when needed |
+| `scheduler_step_maybe_with_cfg()` | Step scheduler with PP sync   | Runs scheduler on the last PP rank, returns updated latents to rank 0           |
+| `_sync_pp_send()`                 | Flush pending async sends     | Waits for outstanding `isend` handles before later collectives or decode        |
+
+`PipelineParallelMixin` is intentionally a pipeline-level abstraction. Your model-specific `predict_noise()` still
+defines how a local stage executes.
+
+### How It Works
+
+`predict_noise_maybe_with_cfg()` automatically switches between these modes:
+
+- **PP disabled** (`pipeline_parallel_size == 1`):
+    - Falls back to `CFGParallelMixin.predict_noise_maybe_with_cfg()`
+- **PP only** (`pipeline_parallel_size > 1`, `cfg_parallel_size == 1`):
+    - Rank 0 starts with the input latents
+    - Middle ranks receive `intermediate_tensors`, run their local layer range, and asynchronously send downstream
+    - The last rank returns the final noise prediction
+    - When CFG is enabled in this mode (sequential CFG), both the positive and negative branches run through the same
+      PP pipeline, doubling the communication volume per denoising step. Prefer PP + CFG-Parallel when
+      `cfg_parallel_size > 1` is available to avoid this overhead.
+- **PP + CFG-Parallel** (`pipeline_parallel_size > 1`, `cfg_parallel_size > 1`):
+    - Each PP pipeline carries one CFG branch
+    - The last PP rank all-gathers across the CFG group
+    - CFG combination happens on every last PP rank in the CFG group, matching the non-PP CFG-parallel behavior
+
+`scheduler_step_maybe_with_cfg()` keeps the denoising loop consistent:
+
+- **PP disabled**:
+    - Falls back to `scheduler_step_maybe_with_cfg()`
+- **PP enabled**:
+    - Only the last PP rank has `noise_pred` and runs the scheduler step
+    - The resulting latents are sent back to rank 0
+    - Rank 0 receives an `AsyncLatents` wrapper that resolves only when the tensor is next consumed
+
+This asynchronous design avoids unnecessary blocking between denoising steps.
+
+When a pipeline class is defined, `PipelineParallelMixin` wraps its `diffuse()` method and calls `_sync_pp_send()` in a
+`finally` block after `diffuse()` returns or raises. This keeps model code free of explicit PP cleanup while still
+ensuring that the final non-blocking PP send completes before decode or later collectives.
+
+---
+
+## Step-by-Step Implementation
+
+### Step 1: Inherit `PipelineParallelMixin` and `CFGParallelMixin`
+
+`PipelineParallelMixin` **requires** `CFGParallelMixin` and must appear before it in the class MRO. This is enforced at
+class definition time via `__init_subclass__`: defining a pipeline that inherits `PipelineParallelMixin` without
+`CFGParallelMixin`, or after `CFGParallelMixin`, raises a `TypeError` immediately on import.
+
+`PipelineParallelMixin` delegates noise prediction, CFG combination, and scheduler stepping to `CFGParallelMixin`, which
+supplies `predict_noise()`, `predict_noise_maybe_with_cfg()`, `scheduler_step_maybe_with_cfg()`, and
+`combine_cfg_noise()`.
+
+**Example:**
+
+```python
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.pipeline_parallel import PipelineParallelMixin
+import torch.nn as nn
+
+
+class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin):
+    ...
+```
+
+The order matters: `PipelineParallelMixin` must be listed before `CFGParallelMixin` so calls to
+`predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()` resolve to the PP-aware wrappers, while their
+`super()` calls delegate to the CFG implementation when PP is disabled or after the last PP stage.
+
+### Step 2: Make model forward and `predict_noise()` PP-aware
+
+The PP mixin injects `intermediate_tensors` into the normal `predict_noise()` call. Your model forward path must support
+two inputs:
+
+- Normal input from rank 0, typically passed as `hidden_states` or `x`
+- `intermediate_tensors` from upstream PP ranks
+
+The standard model forward pattern is:
+
+1. If `intermediate_tensors` is present, read the local hidden state from it.
+2. Run only this rank's layer slice.
+3. Return `IntermediateTensors(...)` on non-last PP ranks.
+4. Return the final model output on the last PP rank. `CFGParallelMixin.predict_noise()` already follows this contract
+   for common pipelines.
+
+**Minimal example:**
+
+```python
+from vllm.sequence import IntermediateTensors
+from vllm_omni.diffusion.distributed.parallel_state import get_pp_group
+
+
+def forward(self, hidden_states=None, intermediate_tensors=None, **kwargs):
+    if intermediate_tensors is not None:
+        hidden_states = intermediate_tensors["hidden_states"]
+
+    for i in range(self.start_layer, self.end_layer):
+        hidden_states = self.layers[i](hidden_states)
+
+    pp_group = get_pp_group()
+    if not pp_group.is_last_rank:
+        return IntermediateTensors({"hidden_states": hidden_states})
+    return (hidden_states,)
+```
+
+### Step 3: Partition the transformer layers
+
+The local module on each PP rank must expose only that rank's layer slice. In the reference tests and vLLM model
+implementations, this is typically done with vLLM utilities such as `make_layers(...)`, which fill missing ranges with
+`PPMissingLayer`. Besides partitioning the layers themselves, model authors should also wire
+`make_empty_intermediate_tensors_factory(...)` for intermediate tensor allocation and `is_pp_missing_parameter(...)` for
+PP-aware weight loading.
+
+To prepare a transformer for PP, implement the following pieces in order.
+
+#### 3.1 Split the transformer layers across PP ranks
+
+Each PP rank should own only its local layer range, typically exposed as `[start_layer, end_layer)`. In practice this is
+usually done with `make_layers(...)`, which constructs local layers and fills missing ranges with `PPMissingLayer`.
+
+The goal is:
+
+- every PP rank knows its `[start_layer, end_layer)` range
+- non-local layers do not execute on this rank
+- the forward can resume from incoming `intermediate_tensors`
+
+By default, layers are distributed evenly across PP ranks using `get_pp_indices()`. When the number of layers is not
+divisible by the PP size, remaining layers are assigned to middle partitions to balance compute and memory. You can
+override this with the `VLLM_PP_LAYER_PARTITION` environment variable to specify an exact per-rank layer count:
+
+```bash
+# Example: 40 layers across 4 PP ranks, assigning 8 / 12 / 12 / 8 layers
+export VLLM_PP_LAYER_PARTITION=8,12,12,8
+```
+
+The value must be a comma-separated list of integers whose length equals `pipeline_parallel_size` and whose sum equals
+the total number of transformer layers.
+
+#### 3.2 Expose `make_empty_intermediate_tensors`
+
+Your transformer module should expose `self.make_empty_intermediate_tensors`, typically created with
+`make_empty_intermediate_tensors_factory(...)`.
+
+This is important because PP ranks need a consistent way to allocate placeholder `IntermediateTensors` with the expected
+keys and hidden dimension.
+
+**Example:**
+
+```python
+from vllm.model_executor.models.utils import make_empty_intermediate_tensors_factory
+
+self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+    ["hidden_states"],
+    inner_dim,
+)
+```
+
+For Wan2.2, the intermediate payload between PP stages is the token sequence stored under `"hidden_states"`, so the
+factory is created with that key and the transformer hidden size.
+
+#### 3.3 Return `IntermediateTensors` on non-last PP ranks
+
+Your model `forward()` or custom `predict_noise()` implementation should consume `intermediate_tensors` on non-first
+ranks and return `IntermediateTensors(...)` on non-last ranks.
+
+This lets each PP stage resume from the upstream hidden states and pass the local result to the next stage.
+
+#### 3.4 Skip non-local weights during `load_weights()`
+
+When a model is PP-partitioned, many parameters in the checkpoint belong to layers that do not exist on the current
+rank. `load_weights()` must therefore skip parameters for missing PP stages using `is_pp_missing_parameter(...)`.
+
+If this is not done, weight loading will either fail or incorrectly try to load tensors into `PPMissingLayer`
+placeholders.
+
+The Wan2.2 transformer is the best reference here: it uses `is_pp_missing_parameter(...)` before loading both remapped
+and fused parameters.
+
+If your model has multiple transformer variants, PP still works as long as each selected transformer obeys the same
+contract.
+
+### Step 4: Use the standard denoising contract
+
+vLLM-Omni diffusion pipelines already route denoising through `diffuse()`, `predict_noise_maybe_with_cfg()`, and
+`scheduler_step_maybe_with_cfg()`. `PipelineParallelMixin` overrides those standard helpers when
+`pipeline_parallel_size > 1`, so model integrations should not add separate PP-specific helper names or manual post-loop
+synchronization.
+
+In PP mode:
+
+- Non-last PP ranks return `None` from `predict_noise_maybe_with_cfg()` after sending `IntermediateTensors` downstream.
+- The last PP rank returns the final prediction, and with CFG-Parallel every last PP rank in the CFG group receives the
+  combined result.
+- Rank 0 receives `AsyncLatents` from `scheduler_step_maybe_with_cfg()`, which resolves only when the latents are
+  consumed.
+- Pending non-blocking sends are flushed automatically when `diffuse()` exits.
+
+---
+
+## Testing
+
+Use an offline inference script with `pipeline_parallel_size > 1`.
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+--width=1280 \
+--height=704 \
+--guidance-scale=5.0 \
+--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
+--output=t2v_5B_pp2.mp4 \
+--pipeline-parallel-size=2
+```
+
+For PP + CFG-Parallel together:
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+--width=1280 \
+--height=704 \
+--guidance-scale=5.0 \
+--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
+--output=t2v_5B_pp2_cfg2.mp4 \
+--pipeline-parallel-size=2 \
+--cfg-parallel-size=2
+```
+
+**Verify:**
+
+1. The run completes without hangs at the PP stage boundary
+2. Output quality matches the non-PP baseline within normal numerical variance
+3. Peak memory per GPU drops compared with the single-rank model
+4. There are no outstanding communication errors before decode
+
+---
+
+## Troubleshooting
+
+### Issue: `TypeError` on import — `CFGParallelMixin` missing
+
+**Symptoms:** Importing a pipeline that inherits `PipelineParallelMixin` raises:
+
+```
+TypeError: YourPipeline inherits PipelineParallelMixin but not CFGParallelMixin.
+```
+
+or:
+
+```
+TypeError: YourPipeline must inherit PipelineParallelMixin before CFGParallelMixin ...
+```
+
+**Cause:** `PipelineParallelMixin` enforces via `__init_subclass__` that the subclass also inherits `CFGParallelMixin`
+and lists `PipelineParallelMixin` first in the MRO.
+
+**Solution:** Add `CFGParallelMixin` to the pipeline's base classes after `PipelineParallelMixin`:
+
+```python
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.pipeline_parallel import PipelineParallelMixin
+
+
+class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin):
+    ...
+```
+
+### Issue: Non-last PP ranks crash when calling `predict_noise`
+
+**Symptoms:** Shape or missing-input errors on ranks other than the first or last PP rank.
+
+**Cause:** The model forward path assumes direct input tensors and ignores `intermediate_tensors`.
+
+**Solution:** Update the transformer `forward()` or custom `predict_noise()` path to load hidden states from
+`intermediate_tensors` when present.
+
+### Issue: PP outputs differ from single-GPU baseline
+
+**Symptoms:** The PP run finishes but produces numerically inconsistent outputs.
+
+**Causes & Solutions:**
+
+- **Local layer partitioning is wrong**
+    - Verify each rank runs only its `[start_layer, end_layer)` slice
+- **Non-last ranks return a plain tensor instead of `IntermediateTensors`**
+    - Return `IntermediateTensors({...})` until the last PP stage
+- **CFG branch wiring is incorrect**
+    - When CFG is enabled, confirm positive and negative kwargs are passed exactly as in the non-PP path
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Component               | Path                                                       | Notes                                                                                                                   |
+|-------------------------|------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------|
+| `PipelineParallelMixin` | `vllm_omni/diffusion/distributed/pipeline_parallel.py`     | Core PP communication and scheduler helpers                                                                             |
+| `CFGParallelMixin`      | `vllm_omni/diffusion/distributed/cfg_parallel.py`          | Default `predict_noise()` tuple normalization and CFG helper fallback                                                   |
+| Wan2.2 transformer      | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py`  | Reference for layer partitioning, `IntermediateTensors`, `make_empty_intermediate_tensors`, and PP-aware weight loading |
+| Wan2.2 T2V pipeline     | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py`     | Reference PP + CFG integration for text-to-video                                                                        |
+| Wan2.2 I2V pipeline     | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py` | Reference PP + CFG integration for image-to-video                                                                       |
+| PP tests                | `tests/diffusion/distributed/test_pipeline_parallel.py`    | Baseline parity and async communication tests                                                                           |
diff --git a/docs/user_guide/diffusion/frame_interpolation.md b/docs/user_guide/diffusion/frame_interpolation.md
index 349af50c5..ab03363cf 100644
--- a/docs/user_guide/diffusion/frame_interpolation.md
+++ b/docs/user_guide/diffusion/frame_interpolation.md
@@ -28,7 +28,6 @@ Frame interpolation is currently supported for:
 
 - `WanPipeline` (Wan2.2 text-to-video)
 - `WanImageToVideoPipeline`
-- `Wan22TI2VPipeline`
 
 ## Request Parameters
 
diff --git a/docs/user_guide/diffusion/parallelism/overview.md b/docs/user_guide/diffusion/parallelism/overview.md
index 90d0b9660..d482b5ea2 100644
--- a/docs/user_guide/diffusion/parallelism/overview.md
+++ b/docs/user_guide/diffusion/parallelism/overview.md
@@ -4,13 +4,14 @@ This guide covers the parallelism methods in vLLM-Omni for speeding up diffusion
 
 ## Supported Methods
 
-| Method | Description |
-|--------|-------------|
-| **[Tensor Parallelism](tensor_parallel.md)** | Shards DiT weights across GPUs to reduce per-GPU memory |
-| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
-| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation |
-| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
-| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs |
-| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) |
+| Method                                             | Description                                                                                                         |
+|----------------------------------------------------|---------------------------------------------------------------------------------------------------------------------|
+| **[Tensor Parallelism](tensor_parallel.md)**       | Shards DiT weights across GPUs to reduce per-GPU memory                                                             |
+| **[Sequence Parallelism](sequence_parallel.md)**   | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
+| **[CFG-Parallel](cfg_parallel.md)**                | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation                         |
+| **[Pipeline Parallelism](pipeline_parallel.md)**   | Splits the denoising transformer block-wise across sequential GPU stages to reduce per-GPU model memory             |
+| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory                                              |
+| **[HSDP](hsdp.md)**                                | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs              |
+| **[Expert Parallelism](expert_parallel.md)**       | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0)                                          |
 
 See [Supported Models](../../diffusion_features.md#supported-models) for per-model compatibility.
diff --git a/docs/user_guide/diffusion/parallelism/pipeline_parallel.md b/docs/user_guide/diffusion/parallelism/pipeline_parallel.md
new file mode 100644
index 000000000..0827e5efa
--- /dev/null
+++ b/docs/user_guide/diffusion/parallelism/pipeline_parallel.md
@@ -0,0 +1,205 @@
+# Pipeline Parallelism Guide
+
+## Table of Content
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Example Script](#example-script)
+- [Configuration Parameters](#configuration-parameters)
+- [Best Practices](#best-practices)
+- [Troubleshooting](#troubleshooting)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+Pipeline Parallelism splits the denoising transformer block-wise into sequential stages across GPUs. Each rank owns only
+part of the transformer, which reduces per-GPU model memory and enables larger diffusion models to run across multiple
+devices.
+
+It can also be combined with other distributed methods such as CFG-Parallel, Tensor Parallelism, and Sequence
+Parallelism.
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+---
+
+## Quick Start
+
+### Basic Usage
+
+Simplest working example:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(
+    model="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+    parallel_config=DiffusionParallelConfig(
+        pipeline_parallel_size=2,
+    ),
+)
+
+outputs = omni.generate(
+    {"prompt": "A cinematic drone shot over snowy mountains"},
+    OmniDiffusionSamplingParams(
+        num_inference_steps=40,
+        num_frames=81,
+        height=704,
+        width=1280,
+    ),
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use python scripts under:
+
+- `examples/offline_inference/text_to_video/text_to_video.py`
+- `examples/offline_inference/image_to_video/image_to_video.py`
+
+Text-to-video example:
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+--width=1280 \
+--height=704 \
+--guidance-scale=5.0 \
+--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
+--output=t2v_5B_pp2.mp4 \
+--pipeline-parallel-size=2
+```
+
+Pipeline Parallelism can also be combined with CFG-Parallel:
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+--width=1280 \
+--height=704 \
+--guidance-scale=5.0 \
+--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
+--output=t2v_5B_pp2_cfg2.mp4 \
+--pipeline-parallel-size=2 \
+--cfg-parallel-size=2
+```
+
+### Online Serving
+
+Enable Pipeline Parallelism in online serving:
+
+```bash
+# Default PP configuration
+vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni --port 8091 --pipeline-parallel-size 2
+
+# PP + CFG-Parallel
+vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni --port 8091 \
+  --pipeline-parallel-size 2 \
+  --cfg-parallel-size 2
+```
+
+---
+
+## Configuration Parameters
+
+In `DiffusionParallelConfig`
+
+| Parameter                | Type | Default | Description                                                                                                      |
+|--------------------------|------|---------|------------------------------------------------------------------------------------------------------------------|
+| `pipeline_parallel_size` | int  | 1       | Number of pipeline-parallel stages. Set to a value greater than 1 to split the denoising transformer across GPUs |
+
+> [!NOTE]
+> Total GPU count is the product of all enabled distributed dimensions, for example
+`pipeline_parallel_size * cfg_parallel_size * tensor_parallel_size * ulysses_degree * ring_degree`.
+
+### Manual Layer Partitioning
+
+By default, transformer layers are distributed evenly across PP ranks. You can override this with the
+`VLLM_PP_LAYER_PARTITION` environment variable to assign a specific number of layers to each rank:
+
+```bash
+# Example: 40 layers across 4 PP ranks, assigning 8 / 12 / 12 / 8 layers
+export VLLM_PP_LAYER_PARTITION=8,12,12,8
+```
+
+The value must be a comma-separated list of integers whose length equals `pipeline_parallel_size` and whose sum equals
+the total number of transformer layers. This is useful when you want to balance memory or compute asymmetrically across
+ranks.
+
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- Large diffusion transformers that do not fit comfortably on one GPU
+- Multi-GPU setups where reducing per-GPU model memory is more important than minimizing communication
+- Combining with CFG-Parallel or other distributed methods on supported models
+
+**Not for:**
+
+- Single GPU setups
+- Models that do not support Pipeline Parallelism
+  (check [supported models](../../diffusion_features.md#supported-models))
+- Very small models where inter-stage communication overhead may outweigh the benefit
+
+### Expected Behavior
+
+Pipeline Parallelism primarily reduces per-GPU memory usage by splitting the transformer into stage-local blocks.
+Depending on the model, topology, and resolution, it may also help execution fit into available hardware, but it is not
+primarily a latency optimization.
+
+---
+
+## Troubleshooting
+
+### Common Issue 1: No benefit from Pipeline Parallelism
+
+**Symptoms**: PP is enabled, but latency does not improve or becomes slightly worse.
+
+**Solutions**:
+
+1. **Check your goal:**
+
+```python
+# PP is mainly for memory scaling, not guaranteed latency speedup
+parallel_config = DiffusionParallelConfig(pipeline_parallel_size=2)
+```
+
+2. **Check model support:**
+    - Verify your model in [supported models](../../diffusion_features.md#supported-models)
+    - PP is currently validated only on selected pipelines
+
+3. **Combine with other methods when appropriate:**
+    - PP can be combined with CFG-Parallel, Tensor Parallelism, or Sequence Parallelism on supported models
+
+### Common Issue 2: PP pipeline fails at import
+
+**Symptoms**: Importing a custom pipeline raises a `TypeError` about `CFGParallelMixin` or mixin order.
+
+**Solutions**:
+
+1. Inherit both mixins.
+2. Put `PipelineParallelMixin` before `CFGParallelMixin` in the class base list.
+3. Use `class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin): ...` as the reference pattern.
+
+---
+
+## Summary
+
+1. ✅ **Enable Pipeline Parallelism** - Set `pipeline_parallel_size > 1` in `DiffusionParallelConfig`
+2. ✅ **Use Supported Models** - Verify your model supports PP in
+   [supported models](../../diffusion_features.md#supported-models)
+3. ✅ **Combine When Needed** - PP can be combined with CFG-Parallel and other distributed methods on supported pipelines
+4. ✅ **Scale for Memory** - Use PP primarily to reduce per-GPU model memory and fit larger transformers
diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md
index 94238f26f..26b4d9bf7 100644
--- a/docs/user_guide/diffusion_features.md
+++ b/docs/user_guide/diffusion_features.md
@@ -35,14 +35,15 @@ Cache methods trade minimal quality for significant speedup. Quality loss is typ
 
 Parallelism methods distribute computation across GPUs without quality loss (mathematically equivalent to single-GPU).
 
-| Method | Description | Best For |
-|--------|-------------|----------|
-| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via all-to-all communication | High-resolution images (>1536px) or long videos with 2-8 GPUs |
-| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via ring-based communication | Videos, very long sequences, memory-constrained, with 2-8 GPUs |
-| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)** | Splits CFG positive/negative branches across devices | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs |
-| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs |
-| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
-| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
+| Method                                                                 | Description                                                              | Best For                                                          |
+|------------------------------------------------------------------------|--------------------------------------------------------------------------|-------------------------------------------------------------------|
+| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)**           | Sequence parallelism via all-to-all communication                        | High-resolution images (>1536px) or long videos with 2-8 GPUs     |
+| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)**       | Sequence parallelism via ring-based communication                        | Videos, very long sequences, memory-constrained, with 2-8 GPUs    |
+| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)**              | Splits CFG positive/negative branches across devices                     | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs    |
+| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)**     | Shards model weights across devices                                      | Large models that don't fit in single GPU, with 2+ GPUs           |
+| **[Pipeline Parallelism](diffusion/parallelism/pipeline_parallel.md)** | Splits the denoising transformer block-wise across sequential GPU stages | Large diffusion transformers that need lower per-GPU model memory |
+| **[HSDP](diffusion/parallelism/hsdp.md)**                              | Weight sharding via FSDP2, redistributed on-demand at runtime            | Very large models (14B+) on limited VRAM, combinable with SP      |
+| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)**     | Shards MoE expert MLP blocks across devices                              | MoE diffusion models (e.g., HunyuanImage3.0)                      |
 
 #### Startup Optimization
 
@@ -105,32 +106,32 @@ The following tables show which models support each feature:
 
 ### ImageGen
 
-| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
-|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
-| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅  | ❌ | ❌ | ❌ |
-| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
-| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
-| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
-| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
-| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **OmniGen2** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
-| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
-| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
-| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
-| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
-| **SenseNova-U1** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
-| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
-| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
+| Model                    | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
+|--------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
+| **Bagel**                |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **FLUX.1-dev**           |     ✅     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
+| **FLUX.1-schnell**       |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
+| **FLUX.2-klein**         |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
+| **FLUX.1-Kontext-dev**   |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **FLUX.2-dev**           |     ✅     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **GLM-Image**            |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ❌         |
+| **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **MagiHuman**            |     ❌     |     ❌      |           ❌           |       ❓        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **MammothModa2(T2I)**    |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **OmniGen2**             |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **Qwen-Image**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |
+| **Qwen-Image-2512**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |
+| **Qwen-Image-Edit**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **Qwen-Image-Edit-2509** |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |        ✅ (decode)         |          ✅           |       ❌        |        ❌         |
+| **Qwen-Image-Layered**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **SenseNova-U1**         |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **Stable-Diffusion3.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **Z-Image**              |     ✅     |     ✅      |           ✅           |       ❓        |   ✅ (TP=2 only)   |          ❌          |   ✅    |             ❌             |      ✅ (decode)      |       ✅        |        ❌         |
+| **ERNIE-Image**          |     ❌     |     ✅      |           ✅           |       ❓        |         ✅         |          ❌          |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 
 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
@@ -138,15 +139,15 @@ The following tables show which models support each feature:
 
 ### VideoGen
 
-| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
-|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
-| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
-| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
-| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
-| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+| Model                        | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
+|------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
+| **Wan2.2**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ✅         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
+| **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ❌         |
+| **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 
 **Frame Interpolation Support**
 
@@ -155,9 +156,9 @@ The following tables show which models support each feature:
 
 ### AudioGen
 
-| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
-|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
-| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ |
+| Model                 | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
+|-----------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
+| **Stable-Audio-Open** |     ✅     |     ❌      |           ❓           |       ❓        |         ❌         |          ❌          |   ✅    |             ✅             |          ❌           |       ✅        |        ❌         |
 
 
 ## Feature Compatibility
@@ -257,7 +258,7 @@ Measured on NVIDIA H800:
 
 **Parallelism Methods:**
 
-- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, HSDP, and Expert Parallelism
+- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, Pipeline Parallelism, HSDP, and Expert Parallelism
 
 **Memory Optimization:**
 
diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py
index 4964f10c4..9e2521233 100644
--- a/examples/offline_inference/image_to_video/image_to_video.py
+++ b/examples/offline_inference/image_to_video/image_to_video.py
@@ -142,6 +142,42 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument(
+        "--enforce-eager",
+        action="store_true",
+        help="Disable torch.compile and force eager execution.",
+    )
+    parser.add_argument(
+        "--audio-sample-rate",
+        type=int,
+        default=24000,
+        help="Sample rate for audio output when saved (default: 24000).",
+    )
+    parser.add_argument(
+        "--cache-backend",
+        type=str,
+        default=None,
+        choices=["cache_dit", "tea_cache"],
+        help=(
+            "Cache backend to use for acceleration. "
+            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
+            "Default: None (no cache acceleration)."
+        ),
+    )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        choices=["fp8", "mxfp8", "int8", "gguf"],
+        help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
+    )
+
+    # Distributed and parallel execution
     parser.add_argument(
         "--ulysses-degree",
         type=int,
@@ -173,28 +209,6 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for VAE patch/tile parallelism (decode).",
     )
-    parser.add_argument(
-        "--enforce-eager",
-        action="store_true",
-        help="Disable torch.compile and force eager execution.",
-    )
-    parser.add_argument(
-        "--audio-sample-rate",
-        type=int,
-        default=24000,
-        help="Sample rate for audio output when saved (default: 24000).",
-    )
-    parser.add_argument(
-        "--cache-backend",
-        type=str,
-        default=None,
-        choices=["cache_dit", "tea_cache"],
-        help=(
-            "Cache backend to use for acceleration. "
-            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
-            "Default: None (no cache acceleration)."
-        ),
-    )
     parser.add_argument(
         "--use-hsdp",
         action="store_true",
@@ -219,16 +233,10 @@ def parse_args() -> argparse.Namespace:
         ),
     )
     parser.add_argument(
-        "--quantization",
-        type=str,
-        default=None,
-        choices=["fp8", "mxfp8", "int8", "gguf"],
-        help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
-    )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
+        "--pipeline-parallel-size",
+        type=int,
+        default=1,
+        help="Number of pipeline parallel stages.",
     )
     parser.add_argument(
         "--profiler-config",
@@ -326,6 +334,7 @@ def main():
         use_hsdp=args.use_hsdp,
         hsdp_shard_size=args.hsdp_shard_size,
         hsdp_replicate_size=args.hsdp_replicate_size,
+        pipeline_parallel_size=args.pipeline_parallel_size,
     )
     omni_kwargs = dict(
         model=args.model,
@@ -366,7 +375,8 @@ def main():
     print(f"  diffusion_kv_cache_skip_layers(config): {args.diffusion_kv_cache_skip_layers}")
     print(
         f"  Parallel configuration: cfg_parallel_size={args.cfg_parallel_size},"
-        f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}"
+        f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size},"
+        f" pipeline_parallel_size={args.pipeline_parallel_size}"
     )
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")
diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py
index b19f0095e..9e1775dd1 100644
--- a/examples/offline_inference/text_to_video/text_to_video.py
+++ b/examples/offline_inference/text_to_video/text_to_video.py
@@ -142,6 +142,32 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument(
+        "--audio-sample-rate",
+        type=int,
+        default=24000,
+        help="Sample rate for audio output when saved (default: 24000).",
+    )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
+    parser.add_argument(
+        "--profiler-config",
+        type=parse_profiler_config,
+        default=None,
+        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
+    )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        choices=["fp8", "mxfp8", "int8", "gguf"],
+        help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
+    )
+
+    # Distributed and parallel execution
     parser.add_argument(
         "--ulysses-mode",
         type=str,
@@ -174,41 +200,23 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
     )
-    parser.add_argument(
-        "--audio-sample-rate",
-        type=int,
-        default=24000,
-        help="Sample rate for audio output when saved (default: 24000).",
-    )
     parser.add_argument(
         "--vae-patch-parallel-size",
         type=int,
         default=1,
         help="Number of GPUs used for VAE patch/tile parallelism (decode).",
     )
+    parser.add_argument(
+        "--pipeline-parallel-size",
+        type=int,
+        default=1,
+        help="Number of pipeline parallel stages.",
+    )
     parser.add_argument(
         "--enable-expert-parallel",
         action="store_true",
         help="Enable expert parallelism for MoE layers.",
     )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
-    parser.add_argument(
-        "--profiler-config",
-        type=parse_profiler_config,
-        default=None,
-        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
-    )
-    parser.add_argument(
-        "--quantization",
-        type=str,
-        default=None,
-        choices=["fp8", "mxfp8", "int8", "gguf"],
-        help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
-    )
     parser.add_argument(
         "--use-hsdp",
         action="store_true",
@@ -262,6 +270,7 @@ def main():
         cfg_parallel_size=args.cfg_parallel_size,
         tensor_parallel_size=args.tensor_parallel_size,
         vae_patch_parallel_size=args.vae_patch_parallel_size,
+        pipeline_parallel_size=args.pipeline_parallel_size,
         enable_expert_parallel=args.enable_expert_parallel,
     )
 
@@ -307,7 +316,8 @@ def main():
     print(
         f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree},"
         f" cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size},"
-        f" vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}"
+        f" vae_patch_parallel_size={args.vae_patch_parallel_size}, pipeline_parallel_size={args.pipeline_parallel_size},"
+        f" enable_expert_parallel={args.enable_expert_parallel}"
     )
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")
diff --git a/tests/diffusion/distributed/test_pipeline_parallel.py b/tests/diffusion/distributed/test_pipeline_parallel.py
new file mode 100644
index 000000000..b76bf8c89
--- /dev/null
+++ b/tests/diffusion/distributed/test_pipeline_parallel.py
@@ -0,0 +1,802 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import socket
+
+import pytest
+import torch
+from vllm.model_executor.models.utils import PPMissingLayer, make_empty_intermediate_tensors_factory, make_layers
+from vllm.sequence import IntermediateTensors
+from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
+
+import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.parallel_state import (
+    destroy_distributed_env,
+    get_classifier_free_guidance_rank,
+    get_pp_group,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.parallel]
+
+
+def _find_free_port() -> str:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return str(s.getsockname()[1])
+
+
+def update_environment_variables(envs_dict: dict[str, str]) -> None:
+    for k, v in envs_dict.items():
+        os.environ[k] = v
+
+
+# ---------------------------------------------------------------------------
+# Shared stubs used by both unit and distributed tests
+# ---------------------------------------------------------------------------
+
+
+class FakeWork:
+    """Drop-in for torch.distributed.Work that records whether wait() was called."""
+
+    def __init__(self):
+        self.waited = False
+
+    def wait(self):
+        self.waited = True
+
+
+class SimpleScheduler:
+    """Minimal diffusion-step scheduler: latents -= 0.1 * noise_pred."""
+
+    def step(self, noise_pred: torch.Tensor, t, latents: torch.Tensor, return_dict: bool = False):
+        return (latents - 0.1 * noise_pred,)
+
+
+class FakeVAE:
+    def __init__(self, distributed_enabled: bool = False):
+        self.calls = 0
+        self.distributed_enabled = distributed_enabled
+
+    def decode(self, z: torch.Tensor):
+        """Original decode docstring."""
+        self.calls += 1
+        return (z + 1,)
+
+    def is_distributed_enabled(self) -> bool:
+        return self.distributed_enabled
+
+
+class MockPipelineParallel(PipelineParallelMixin, CFGParallelMixin):
+    """Minimal pipeline used to exercise PipelineParallelMixin.
+
+    Uses vLLM's ``make_layers`` for layer partitioning — the same utility used
+    by real DiT models — so the PP layer-split logic is exercised faithfully.
+
+    Each layer's weights are seeded by ``seed + layer_index`` so that layer ``i``
+    is initialized identically on every rank regardless of which ranks are active,
+    allowing the distributed output to be compared against the single-GPU baseline.
+
+    Args:
+        num_layers: Total number of Linear layers.
+        dim:        Input / hidden dimension.
+        seed:       Base RNG seed; layer ``i`` uses ``seed + i``.
+        device:     Target device for layer weights (default: CPU).
+        dtype:      Target dtype for layer weights (default: float32).
+    """
+
+    def __init__(
+        self,
+        num_layers: int = 4,
+        dim: int = 64,
+        seed: int = 42,
+        device: torch.device | None = None,
+        dtype: torch.dtype = torch.float32,
+    ):
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            num_layers,
+            lambda prefix: torch.nn.Linear(dim, dim, bias=False),
+            prefix="layers",
+        )
+
+        for i, layer in enumerate(self.layers):
+            if not isinstance(layer, PPMissingLayer):
+                torch.manual_seed(seed + i)
+                torch.nn.init.normal_(layer.weight, mean=0.0, std=0.02)
+
+        self.layers.to(device=device, dtype=dtype)
+
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(["hidden_states"], dim)
+        self.scheduler = SimpleScheduler()
+
+    def predict_noise(self, x=None, intermediate_tensors=None, **_kwargs) -> torch.Tensor | IntermediateTensors:
+        """Layer-split forward pass.
+
+        * First PP rank: uses ``x`` from caller kwargs.
+        * Later PP ranks: overrides ``x`` with ``intermediate_tensors["hidden_states"]``
+          (which transparently waits for the async receive).
+        * Non-last PP ranks return ``IntermediateTensors``; the last rank
+          returns the plain noise-prediction tensor.
+        """
+        if intermediate_tensors is not None:
+            x = intermediate_tensors["hidden_states"]
+
+        for i in range(self.start_layer, self.end_layer):
+            x = self.layers[i](x)
+
+        pp_group = get_pp_group()
+        if not pp_group.is_last_rank:
+            return IntermediateTensors({"hidden_states": x})
+        return x
+
+
+# ---------------------------------------------------------------------------
+# 1.  AsyncLatents – unit tests (no distributed env required)
+# ---------------------------------------------------------------------------
+
+
+class TestAsyncLatents:
+    """Verifies the lazy-resolution behaviour of AsyncLatents without a real process group."""
+
+    pytestmark = [pytest.mark.cpu]
+
+    def _make(self, tensor: torch.Tensor, handles: list | None = None, postproc: list | None = None) -> AsyncLatents:
+        return AsyncLatents({"latents": tensor}, handles or [], postproc or [])
+
+    def test_resolve_returns_wrapped_tensor(self):
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        assert al._resolve() is t
+
+    def test_attribute_access_resolves(self):
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        assert al.shape == t.shape
+        assert al.dtype == t.dtype
+
+    def test_torch_function_protocol(self):
+        """torch ops that receive an AsyncLatents should see the underlying tensor."""
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        mask = torch.ones_like(t)
+        result = mask * al  # triggers __torch_function__
+        torch.testing.assert_close(result, mask * t)
+
+    def test_torch_function_with_list_arg(self):
+        """__torch_function__ must unwrap AsyncLatents inside list/tuple args."""
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        result = torch.cat([al, al], dim=0)
+        torch.testing.assert_close(result, torch.cat([t, t], dim=0))
+
+    def test_torch_tensor_conversion(self):
+        """torch.as_tensor on an AsyncLatents must share storage with the underlying tensor (no copy)."""
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        result = torch.as_tensor(al)
+        assert result.data_ptr() == t.data_ptr(), "torch.as_tensor copied the data instead of sharing storage"
+
+    def test_handles_are_waited_before_resolve(self):
+        t = torch.randn(2, 4)
+        h1, h2 = FakeWork(), FakeWork()
+        al = self._make(t, handles=[h1, h2])
+        _ = al.shape  # trigger resolution
+        assert h1.waited and h2.waited, "Not all handles were waited on"
+
+    def test_postproc_callbacks_invoked_on_resolve(self):
+        t = torch.randn(2, 4)
+        log: list[int] = []
+        al = self._make(t, postproc=[lambda: log.append(1), lambda: log.append(2)])
+        _ = al.shape
+        assert log == [1, 2], f"postproc not called in order: {log}"
+
+    def test_idempotent_resolve(self):
+        """handle.wait() must not be called twice if _resolve() is called twice."""
+        t = torch.randn(2, 4)
+        h = FakeWork()
+        al = self._make(t, handles=[h])
+        _ = al.shape  # first resolve
+        h.waited = False  # reset sentinel
+        _ = al.dtype  # second resolve
+        assert not h.waited, "handle.wait() was called a second time"
+
+
+# ---------------------------------------------------------------------------
+# 2.  _sync_pp_send / diffuse wrapper – unit tests (no distributed env required)
+# ---------------------------------------------------------------------------
+
+
+class TestSyncPPSend:
+    """Verifies PipelineParallelMixin's internal PP-send flush."""
+
+    pytestmark = [pytest.mark.cpu]
+
+    @staticmethod
+    def _make_pipeline() -> PipelineParallelMixin:
+        # Instantiate a bare mixin — no layers, no distributed env needed.
+        # _sync_pp_send only touches _pp_send_work, so this is sufficient.
+        class _BarePP(PipelineParallelMixin, CFGParallelMixin):
+            pass
+
+        return _BarePP()
+
+    def test_noop_when_work_list_empty(self):
+        pipeline = self._make_pipeline()
+        pipeline._sync_pp_send()
+        assert pipeline._pp_send_work == []
+
+    def test_waits_all_pending_handles(self):
+        pipeline = self._make_pipeline()
+        works = [FakeWork(), FakeWork(), FakeWork()]
+        pipeline._pp_send_work = works
+        pipeline._sync_pp_send()
+        assert all(w.waited for w in works), "Some handles were not waited on"
+
+    def test_clears_work_list_after_sync(self):
+        pipeline = self._make_pipeline()
+        pipeline._pp_send_work = [FakeWork()]
+        pipeline._sync_pp_send()
+        assert pipeline._pp_send_work == []
+
+
+class TestDiffuseWrapper:
+    """Verifies that PipelineParallelMixin flushes pending sends when diffuse() exits."""
+
+    pytestmark = [pytest.mark.cpu]
+
+    def test_diffuse_flushes_pending_sends_on_success(self):
+        work = FakeWork()
+
+        class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
+            def diffuse(self):
+                self._pp_send_work = [work]
+                return "done"
+
+        pipeline = _DiffusePP()
+
+        assert pipeline.diffuse() == "done"
+        assert work.waited
+        assert pipeline._pp_send_work == []
+
+    def test_diffuse_flushes_pending_sends_on_exception(self):
+        work = FakeWork()
+
+        class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
+            def diffuse(self):
+                self._pp_send_work = [work]
+                raise RuntimeError("boom")
+
+        pipeline = _DiffusePP()
+
+        with pytest.raises(RuntimeError, match="boom"):
+            pipeline.diffuse()
+        assert work.waited
+        assert pipeline._pp_send_work == []
+
+    def test_diffuse_wrapper_preserves_metadata(self):
+        class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
+            def diffuse(self):
+                """Original diffuse docstring."""
+                return "done"
+
+        assert _DiffusePP.diffuse.__name__ == "diffuse"
+        assert _DiffusePP.diffuse.__doc__ == "Original diffuse docstring."
+
+
+class TestVaeDecodeGuard:
+    pytestmark = [pytest.mark.cpu]
+
+    @staticmethod
+    def _make_pipeline(distributed_enabled: bool = False) -> PipelineParallelMixin:
+        class _DecodePP(PipelineParallelMixin, CFGParallelMixin):
+            def __init__(self):
+                self.vae = FakeVAE(distributed_enabled=distributed_enabled)
+
+        return _DecodePP()
+
+    @staticmethod
+    def _set_rank(monkeypatch, world_size: int, first_stage: bool) -> None:
+        monkeypatch.setattr(pp_module, "get_pipeline_parallel_world_size", lambda: world_size)
+        monkeypatch.setattr(pp_module, "is_pipeline_first_stage", lambda: first_stage)
+
+    def test_calls_original_decode_when_pp_disabled(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=1, first_stage=True)
+        pipeline = self._make_pipeline()
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)[0]
+
+        assert pipeline.vae.calls == 1
+        torch.testing.assert_close(output, z + 1)
+
+    def test_calls_original_decode_on_first_stage(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=2, first_stage=True)
+        pipeline = self._make_pipeline()
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)[0]
+
+        assert pipeline.vae.calls == 1
+        torch.testing.assert_close(output, z + 1)
+
+    def test_skips_decode_on_non_first_stage(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=2, first_stage=False)
+        pipeline = self._make_pipeline()
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)
+
+        assert pipeline.vae.calls == 0
+        assert output == (None,)
+
+    def test_calls_original_decode_when_distributed_vae_enabled(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=2, first_stage=False)
+        pipeline = self._make_pipeline(distributed_enabled=True)
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)[0]
+
+        assert pipeline.vae.calls == 1
+        torch.testing.assert_close(output, z + 1)
+
+    def test_decode_wrapper_preserves_metadata(self):
+        pipeline = self._make_pipeline()
+
+        assert pipeline.vae.decode.__name__ == "decode"
+        assert pipeline.vae.decode.__doc__ == "Original decode docstring."
+
+
+@pytest.mark.cpu
+def test_pipeline_parallel_requires_cfg_mixin():
+    with pytest.raises(TypeError, match="inherits PipelineParallelMixin but not CFGParallelMixin"):
+
+        class _MissingCFG(PipelineParallelMixin):
+            pass
+
+
+@pytest.mark.cpu
+def test_pipeline_parallel_requires_mro_before_cfg_mixin():
+    with pytest.raises(TypeError, match="must inherit PipelineParallelMixin before CFGParallelMixin"):
+
+        class _WrongOrder(CFGParallelMixin, PipelineParallelMixin):
+            pass
+
+
+# ---------------------------------------------------------------------------
+# Distributed test helpers
+# ---------------------------------------------------------------------------
+
+
+def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.device:
+    """Initialise the distributed environment for a spawned worker."""
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
+    update_environment_variables(
+        {
+            "RANK": str(local_rank),
+            "LOCAL_RANK": str(local_rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": master_port,
+        }
+    )
+    init_distributed_environment()
+    return device
+
+
+def make_pipeline_and_inputs(
+    test_config: dict, dtype: torch.dtype, device: torch.device, do_true_cfg: bool = False
+) -> tuple["MockPipelineParallel", dict, dict | None]:
+    """Create a MockPipelineParallel and seeded inputs from a test_config dict.
+
+    Must be called after ``initialize_model_parallel`` so that ``make_layers``
+    can read the PP group to determine this rank's layer slice.
+
+    Returns ``(pipeline, positive_kwargs, negative_kwargs)``.
+    ``negative_kwargs`` is ``None`` when ``do_true_cfg=False``.
+    """
+    pipeline = MockPipelineParallel(
+        num_layers=test_config["num_layers"],
+        dim=test_config["dim"],
+        seed=test_config["model_seed"],
+        device=device,
+        dtype=dtype,
+    )
+
+    torch.manual_seed(test_config["input_seed"])
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(test_config["input_seed"])
+    pos_x = {"x": torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)}
+
+    negative_kwargs = None
+    if do_true_cfg:
+        torch.manual_seed(test_config["input_seed"] + 1)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(test_config["input_seed"] + 1)
+        neg_x = torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)
+        negative_kwargs = {"x": neg_x}
+
+    return pipeline, pos_x, negative_kwargs
+
+
+# ---------------------------------------------------------------------------
+# 3.  isend_tensor_dict / irecv_tensor_dict  (2 GPUs)
+# ---------------------------------------------------------------------------
+
+
+def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, result_queue):
+    device = init_dist(local_rank, world_size, master_port)
+    initialize_model_parallel(pipeline_parallel_size=world_size)
+    pp_group = get_pp_group()
+
+    if pp_group.is_first_rank:
+        torch.manual_seed(77)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(77)
+        tensor = torch.randn(3, 5, dtype=torch.float32, device=device)
+        handles = pp_group.isend_tensor_dict({"t": tensor})
+        for h in handles:
+            h.wait()
+        result_queue.put(("sent", tensor.cpu()))
+    else:
+        received = AsyncIntermediateTensors(*pp_group.irecv_tensor_dict())
+        result_queue.put(("received", received["t"].cpu()))
+
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    destroy_distributed_env()
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs")
+@pytest.mark.parametrize("pp_size", [2])
+def test_isend_irecv_tensor_dict(pp_size: int):
+    """isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss."""
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    q = manager.Queue()
+
+    port = _find_free_port()
+    torch.multiprocessing.spawn(isend_irecv_worker, args=(pp_size, port, q), nprocs=pp_size)
+
+    results = {label: tensor for label, tensor in [q.get(), q.get()]}
+    torch.testing.assert_close(
+        results["received"], results["sent"], rtol=0, atol=0, msg="isend/irecv transferred tensor incorrectly"
+    )
+
+
+# ---------------------------------------------------------------------------
+# 4.  predict_noise_maybe_with_cfg
+# ---------------------------------------------------------------------------
+
+_baseline_cache: dict[tuple, torch.Tensor] = {}
+
+
+def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_cfg: bool) -> torch.Tensor:
+    """Compute expected single-GPU output using the same MockPipelineParallel.
+
+    Initializes a trivial distributed env (world_size=1) so that ``make_layers`` and the PP/CFG mixins work normally.
+    Results are cached so identical configs are only computed once.
+    """
+    key = (
+        test_config["num_layers"],
+        test_config["dim"],
+        test_config["batch_size"],
+        test_config["model_seed"],
+        test_config["input_seed"],
+        test_config["cfg_scale"],
+        dtype,
+        do_true_cfg,
+    )
+    if key in _baseline_cache:
+        return _baseline_cache[key]
+
+    device = init_dist(0, 1, _find_free_port())
+    initialize_model_parallel(pipeline_parallel_size=1)
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, dtype, device, do_true_cfg=do_true_cfg
+    )
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+
+    destroy_distributed_env()
+
+    _baseline_cache[key] = noise_pred.cpu()
+    return _baseline_cache[key]
+
+
+def predict_noise_worker(
+    local_rank: int,
+    world_size: int,
+    master_port: str,
+    pp_size: int,
+    cfg_size: int,
+    do_true_cfg: bool,
+    dtype: torch.dtype,
+    test_config: dict,
+    result_queue,
+):
+    """Generic predict-noise worker parameterized by PP and CFG topology."""
+    device = init_dist(local_rank, world_size, master_port)
+    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
+
+    pp_group = get_pp_group()
+    cfg_rank = get_classifier_free_guidance_rank()
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, dtype, device, do_true_cfg=do_true_cfg
+    )
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+    # This worker exercises predict_noise_maybe_with_cfg directly, bypassing diffuse().
+    # Flush the non-last PP rank's async send before barrier / process teardown.
+    pipeline._sync_pp_send()
+
+    if pp_group.is_last_rank:
+        assert noise_pred is not None
+        if cfg_rank == 0:
+            result_queue.put(noise_pred.cpu())
+    else:
+        assert noise_pred is None
+
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    destroy_distributed_env()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol",
+    [
+        pytest.param(
+            2,
+            1,
+            False,
+            torch.float32,
+            4,
+            100,
+            1e-5,
+            1e-5,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-no_cfg-float32",
+        ),
+        pytest.param(
+            2,
+            1,
+            False,
+            torch.bfloat16,
+            4,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-no_cfg-bfloat16",
+        ),
+        pytest.param(
+            2,
+            1,
+            True,
+            torch.bfloat16,
+            4,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-seq_cfg-bfloat16",
+        ),
+        pytest.param(
+            2,
+            2,
+            True,
+            torch.bfloat16,
+            4,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
+            id="pp2-cfg2-bfloat16",
+        ),
+        pytest.param(
+            3,
+            1,
+            False,
+            torch.bfloat16,
+            6,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 3, reason="Need at least 3 GPUs"),
+            id="pp3-no_cfg-bfloat16",
+        ),
+    ],
+)
+def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol):
+    """predict_noise_maybe_with_cfg output matches the single-GPU baseline across PP / CFG topologies."""
+    test_config = {
+        "num_layers": num_layers,
+        "dim": 64,
+        "batch_size": 2,
+        "cfg_scale": 7.5,
+        "model_seed": 42,
+        "input_seed": input_seed,
+    }
+
+    baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg)
+
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    pp_q = manager.Queue()
+
+    world_size = pp_size * cfg_size
+    port = _find_free_port()
+    torch.multiprocessing.spawn(
+        predict_noise_worker,
+        args=(world_size, port, pp_size, cfg_size, do_true_cfg, dtype, test_config, pp_q),
+        nprocs=world_size,
+    )
+
+    pp_out = pp_q.get()
+
+    assert baseline_out.shape == pp_out.shape
+    torch.testing.assert_close(
+        pp_out,
+        baseline_out,
+        rtol=rtol,
+        atol=atol,
+        msg=f"PP={pp_size} cfg={cfg_size} {'with' if do_true_cfg else 'no'} CFG output differs from baseline ({dtype=})",
+    )
+
+
+# ---------------------------------------------------------------------------
+# 5.  scheduler_step_maybe_with_cfg
+# ---------------------------------------------------------------------------
+
+
+def compute_scheduler_step_baseline(test_config: dict, do_true_cfg: bool) -> torch.Tensor:
+    """Single-GPU reference: predict_noise + scheduler_step."""
+    device = init_dist(0, 1, _find_free_port())
+    initialize_model_parallel(pipeline_parallel_size=1)
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, torch.float32, device, do_true_cfg=do_true_cfg
+    )
+    latents = positive_kwargs["x"]
+    t = torch.tensor(500, device=device)
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+        result = pipeline.scheduler_step_maybe_with_cfg(
+            noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
+        )
+
+    destroy_distributed_env()
+    return result.cpu()
+
+
+def scheduler_step_worker(
+    local_rank: int,
+    world_size: int,
+    master_port: str,
+    pp_size: int,
+    cfg_size: int,
+    do_true_cfg: bool,
+    test_config: dict,
+    result_queue,
+):
+    device = init_dist(local_rank, world_size, master_port)
+    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
+
+    pp_group = get_pp_group()
+    cfg_rank = get_classifier_free_guidance_rank()
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, torch.float32, device, do_true_cfg=do_true_cfg
+    )
+    latents = positive_kwargs["x"]
+    t = torch.tensor(500, device=device)
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+        latents = pipeline.scheduler_step_maybe_with_cfg(
+            noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
+        )
+    # This worker exercises scheduler_step_maybe_with_cfg directly, bypassing diffuse().
+    # Flush the last PP rank's async latent send before barrier / process teardown.
+    pipeline._sync_pp_send()
+
+    if pp_group.is_first_rank and cfg_rank == 0:
+        resolved = latents.contiguous()
+        result_queue.put(resolved.cpu())
+
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    destroy_distributed_env()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, input_seed",
+    [
+        pytest.param(
+            2,
+            1,
+            False,
+            300,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-no_cfg",
+        ),
+        pytest.param(
+            2,
+            2,
+            True,
+            600,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
+            id="pp2-cfg2-true_cfg",
+        ),
+    ],
+)
+def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
+    """Rank 0 latents after scheduler_step match the single-GPU baseline across PP / CFG topologies."""
+    test_config = {
+        "num_layers": 4,
+        "dim": 64,
+        "batch_size": 2,
+        "cfg_scale": 7.5,
+        "model_seed": 42,
+        "input_seed": input_seed,
+    }
+
+    baseline = compute_scheduler_step_baseline(test_config, do_true_cfg)
+
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    q = manager.Queue()
+
+    port = _find_free_port()
+    world_size = pp_size * cfg_size
+    torch.multiprocessing.spawn(
+        scheduler_step_worker,
+        args=(world_size, port, pp_size, cfg_size, do_true_cfg, test_config, q),
+        nprocs=world_size,
+    )
+
+    result = q.get()
+    torch.testing.assert_close(
+        result,
+        baseline,
+        rtol=0,
+        atol=0,
+        msg=f"PP={pp_size} CFG={cfg_size} scheduler step latents on rank 0 do not match single-GPU baseline",
+    )
diff --git a/tests/diffusion/models/wan2_2/test_wan22_ti2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_ti2v_pipeline.py
deleted file mode 100644
index e611c37b6..000000000
--- a/tests/diffusion/models/wan2_2/test_wan22_ti2v_pipeline.py
+++ /dev/null
@@ -1,97 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-from PIL import Image
-from torch import nn
-
-from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
-from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_ti2v import (
-    Wan22TI2VPipeline,
-    get_wan22_ti2v_pre_process_func,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
-
-
-def _make_ti2v_pipeline() -> Wan22TI2VPipeline:
-    pipeline = object.__new__(Wan22TI2VPipeline)
-    nn.Module.__init__(pipeline)
-    pipeline.device = torch.device("cpu")
-    pipeline.transformer = StubTransformer(in_channels=4, out_channels=4)
-    pipeline.vae = StubVAE(z_dim=4)
-    pipeline.vae_scale_factor_temporal = 4
-    pipeline.vae_scale_factor_spatial = 8
-    pipeline.progress_bar = noop_progress_bar
-    return pipeline
-
-
-def test_ti2v_preprocess_uses_720p_area_for_image_condition() -> None:
-    preprocess = get_wan22_ti2v_pre_process_func(SimpleNamespace())
-    request = SimpleNamespace(
-        prompts=[{"prompt": "p", "multi_modal_data": {"image": Image.new("RGB", (320, 160), "blue")}}],
-        sampling_params=SimpleNamespace(height=None, width=None),
-    )
-
-    result = preprocess(request)
-
-    assert result.sampling_params.height == 672
-    assert result.sampling_params.width == 1344
-    assert result.prompts[0]["multi_modal_data"]["image"].size == (1344, 672)
-
-
-def test_ti2v_diffuse_without_image_condition_expands_patch_timesteps() -> None:
-    pipeline = _make_ti2v_pipeline()
-    latents = torch.zeros(1, 4, 2, 4, 4)
-    calls = []
-
-    def fake_predict_noise_maybe_with_cfg(**kwargs):
-        calls.append(kwargs)
-        return torch.ones_like(latents)
-
-    pipeline.predict_noise_maybe_with_cfg = fake_predict_noise_maybe_with_cfg  # type: ignore[method-assign]
-    pipeline.scheduler_step_maybe_with_cfg = lambda noise, t, current, cfg: current + noise  # type: ignore[method-assign]
-
-    result = pipeline.diffuse(
-        latents=latents,
-        timesteps=torch.tensor([7]),
-        prompt_embeds=torch.zeros(1, 2, 3),
-        negative_prompt_embeds=torch.zeros(1, 2, 3),
-        guidance_scale=3.0,
-        dtype=torch.float32,
-        attention_kwargs={"a": "b"},
-        num_latent_frames=2,
-        latent_height=4,
-        latent_width=4,
-    )
-
-    positive = calls[0]["positive_kwargs"]
-    assert calls[0]["do_true_cfg"] is True
-    assert positive["timestep"].shape == (1, 8)
-    torch.testing.assert_close(positive["timestep"], torch.full((1, 8), 7, dtype=positive["timestep"].dtype))
-    torch.testing.assert_close(positive["hidden_states"], latents)
-    torch.testing.assert_close(result, torch.ones_like(latents))
-
-
-def test_ti2v_prepare_i2v_latents_encodes_condition_and_masks_first_frame() -> None:
-    pipeline = _make_ti2v_pipeline()
-    latents, latent_condition, first_frame_mask = pipeline.prepare_i2v_latents(
-        image=torch.zeros(1, 3, 16, 16),
-        batch_size=1,
-        num_channels_latents=4,
-        height=16,
-        width=16,
-        num_frames=5,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
-        generator=None,
-        latents=torch.zeros(1, 4, 2, 2, 2),
-    )
-
-    torch.testing.assert_close(latents, torch.zeros(1, 4, 2, 2, 2))
-    assert latent_condition.shape == (1, 4, 1, 2, 2)
-    assert first_frame_mask[:, :, 0].sum() == 0
-    assert first_frame_mask[:, :, 1].sum() == 4
diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py
index 5f9047fba..daf9b3939 100644
--- a/vllm_omni/diffusion/cache/cache_dit_backend.py
+++ b/vllm_omni/diffusion/cache/cache_dit_backend.py
@@ -1505,7 +1505,6 @@ CUSTOM_DIT_ENABLERS.update(
     {
         "Wan22Pipeline": enable_cache_for_wan22,
         "Wan22I2VPipeline": enable_cache_for_wan22,
-        "Wan22TI2VPipeline": enable_cache_for_wan22,
         "HunyuanImage3Pipeline": enable_cache_for_hunyuan_image3,
         "FluxPipeline": enable_cache_for_flux,
         "Flux2KleinPipeline": enable_cache_for_flux2_klein,
diff --git a/vllm_omni/diffusion/distributed/cfg_parallel.py b/vllm_omni/diffusion/distributed/cfg_parallel.py
index 98757006b..6810a447b 100644
--- a/vllm_omni/diffusion/distributed/cfg_parallel.py
+++ b/vllm_omni/diffusion/distributed/cfg_parallel.py
@@ -10,6 +10,7 @@ from typing import Any
 
 import torch
 from vllm.logger import init_logger
+from vllm.sequence import IntermediateTensors
 
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_cfg_group,
@@ -369,7 +370,7 @@ class CFGParallelMixin(metaclass=ABCMeta):
             results.append(comb)
         return _unwrap(tuple(results))
 
-    def predict_noise(self, *args: Any, **kwargs: Any) -> torch.Tensor | tuple[torch.Tensor, ...]:
+    def predict_noise(self, *args: Any, **kwargs: Any) -> torch.Tensor | tuple[torch.Tensor, ...] | IntermediateTensors:
         """
         Forward pass through transformer to predict noise.
 
@@ -381,8 +382,11 @@ class CFGParallelMixin(metaclass=ABCMeta):
             multi-output models (e.g., video + audio). Multi-output models
             must also override combine_cfg_noise() and set self.scheduler
             to a composite scheduler that handles tuples.
+            Non-last Pipeline Parallel stages return `IntermediateTensors`
+            instead of final noise tensors wrapped in a tuple.
         """
-        return self.transformer(*args, **kwargs)[0]
+        result = self.transformer(*args, **kwargs)
+        return result if isinstance(result, IntermediateTensors) else result[0]
 
     def diffuse(
         self,
diff --git a/vllm_omni/diffusion/distributed/group_coordinator.py b/vllm_omni/diffusion/distributed/group_coordinator.py
index 5294e6c9e..5c5b092ef 100644
--- a/vllm_omni/diffusion/distributed/group_coordinator.py
+++ b/vllm_omni/diffusion/distributed/group_coordinator.py
@@ -429,6 +429,78 @@ class GroupCoordinator:
                 async_handle.wait()
         return tensor_dict
 
+    def isend_tensor_dict(
+        self,
+        tensor_dict: dict[str, torch.Tensor | Any],
+        dst: int | None = None,
+    ) -> list[torch.distributed.Work]:
+        """Non-blocking send of a tensor dictionary.
+
+        Sends metadata via the Gloo CPU group (blocking) then starts a
+        non-blocking NCCL isend for each GPU tensor.  Returns the list of
+        Work handles; the caller must call handle.wait() before the tensors
+        can be safely reused or freed.
+
+        NOTE: `dst` is the group rank of the destination.
+        """
+        if not torch.distributed.is_initialized() or self.world_size == 1:
+            return []
+
+        if dst is None:
+            dst = self.group_next_rank
+        assert dst < self.world_size, f"Invalid dst rank ({dst})"
+
+        metadata_list, tensor_list = _split_tensor_dict(tensor_dict)
+        self.send_object(metadata_list, dst=dst)
+
+        handles: list[torch.distributed.Work] = []
+        for tensor in tensor_list:
+            if tensor.numel() == 0:
+                continue
+            group = self.cpu_group if tensor.is_cpu else self.device_group
+            handle = torch.distributed.isend(tensor, dst=self.ranks[dst], group=group)
+            if tensor.is_cuda:
+                # Keep allocator from reusing this CUDA buffer before the async send finishes.
+                tensor.record_stream(torch.cuda.current_stream(tensor.device))
+            handles.append(handle)
+        return handles
+
+    def irecv_tensor_dict(
+        self,
+        src: int | None = None,
+    ) -> tuple[dict[str, torch.Tensor | Any], list[torch.distributed.Work], list]:
+        """Non-blocking receive of a tensor dictionary.
+
+        Receives metadata via the Gloo CPU group (blocking) then starts a
+        non-blocking NCCL irecv for each GPU tensor.  Returns
+        ``(tensor_dict, comm_handles, comm_postprocess)`` matching the
+        interface expected by ``AsyncIntermediateTensors``.
+
+        NOTE: `src` is the group rank of the source.
+        """
+        if not torch.distributed.is_initialized() or self.world_size == 1:
+            return {}, [], []
+
+        if src is None:
+            src = self.group_prev_rank
+        assert src < self.world_size, f"Invalid src rank ({src})"
+
+        recv_metadata_list = self.recv_object(src=src)
+        tensor_dict: dict[str, Any] = {}
+        handles: list[torch.distributed.Work] = []
+
+        for key, value in recv_metadata_list:
+            if isinstance(value, TensorMetadata):
+                tensor = torch.empty(value.size, dtype=value.dtype, device=value.device)
+                if tensor.numel() > 0:
+                    group = self.cpu_group if tensor.is_cpu else self.device_group
+                    handles.append(torch.distributed.irecv(tensor, src=self.ranks[src], group=group))
+                _update_nested_dict(tensor_dict, key, tensor)
+            else:
+                _update_nested_dict(tensor_dict, key, value)
+
+        return tensor_dict, handles, []
+
     def send_tensor_dict(
         self,
         tensor_dict: dict[str, torch.Tensor | Any],
diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py
index 82c9c1a45..163f9dd45 100644
--- a/vllm_omni/diffusion/distributed/parallel_state.py
+++ b/vllm_omni/diffusion/distributed/parallel_state.py
@@ -746,11 +746,6 @@ def initialize_model_parallel(
             f" but got {sequence_parallel_size} != {ring_degree} * {ulysses_degree}"
         )
 
-    # FIXME: Since the async p2p communication operation of NPU is not same as cuda in torch,
-    # the pipefusion is not ready for npu yet
-    if current_omni_platform.is_npu():
-        assert pipeline_parallel_size == 1, "Current pipefusion is not ready for NPU"
-
     dit_parallel_size = (
         data_parallel_size * cfg_parallel_size * sequence_parallel_size * pipeline_parallel_size * tensor_parallel_size
     )
diff --git a/vllm_omni/diffusion/distributed/pipeline_parallel.py b/vllm_omni/diffusion/distributed/pipeline_parallel.py
new file mode 100644
index 000000000..b39d8052e
--- /dev/null
+++ b/vllm_omni/diffusion/distributed/pipeline_parallel.py
@@ -0,0 +1,291 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from functools import wraps
+from typing import Any
+
+import torch
+from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
+
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_cfg_group,
+    get_classifier_free_guidance_rank,
+    get_classifier_free_guidance_world_size,
+    get_pipeline_parallel_world_size,
+    get_pp_group,
+    is_pipeline_first_stage,
+)
+
+
+class AsyncLatents:
+    """Transparent async wrapper returned by scheduler_step on rank 0.
+
+    Wraps a pending ``irecv_tensor_dict`` and defers ``handle.wait()`` until the
+    underlying tensor is actually consumed — either via attribute access
+    (e.g. ``latents.to(dtype)``, ``latents.shape``) or via a torch operation
+    (e.g. ``mask * latents``).  This keeps the first PP rank non-blocking after
+    posting the receive, matching the async philosophy used everywhere else in
+    the PP communication layer.
+    """
+
+    __slots__ = ("_tensor_dict", "_handles", "_postproc", "_tensor")
+
+    def __init__(
+        self,
+        tensor_dict: dict[str, torch.Tensor],
+        handles: list[torch.distributed.Work],
+        postproc: list,
+    ):
+        self._tensor_dict = tensor_dict
+        self._handles = handles
+        self._postproc = postproc
+        self._tensor: torch.Tensor | None = None
+
+    def _resolve(self) -> torch.Tensor:
+        if self._tensor is not None:
+            return self._tensor
+        for h in self._handles:
+            h.wait()
+        for fn in self._postproc:
+            fn()
+        self._tensor = self._tensor_dict["latents"]
+        return self._tensor
+
+    # Attribute access (e.g. .shape, .to(), .dtype) delegates to the resolved tensor.
+    def __getattr__(self, name: str):
+        return getattr(self._resolve(), name)
+
+    # Torch function protocol: any torch op involving an AsyncLatents resolves it first.
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
+        kwargs = kwargs or {}
+
+        def _unwrap(x):
+            if isinstance(x, AsyncLatents):
+                return x._resolve()
+            if isinstance(x, (list, tuple)):
+                return type(x)(_unwrap(item) for item in x)  # type(x) return the class of x to preserve its type
+            return x
+
+        args = tuple(_unwrap(a) for a in args)
+        kwargs = {k: _unwrap(v) for k, v in kwargs.items()}
+        return func(*args, **kwargs)
+
+
+class PipelineParallelMixin:
+    """
+    Mixin providing Pipeline Parallelism for diffusion pipelines.
+
+    All PP ranks run the full denoising loop in `forward()`.
+    `predict_noise_maybe_with_cfg` and `scheduler_step_maybe_with_cfg` encapsulate all inter-rank communication.
+
+    Communication pattern per denoising step:
+      Forward chain : rank 0 → 1 → … → N-1  via async isend/irecv (AsyncIntermediateTensors)
+      Next timestep : last rank → rank 0     via async isend/irecv (AsyncLatents)
+
+    All communication is asynchronous using isend_tensor_dict/irecv_tensor_dict.
+    Only rank 0 needs updated latents for the next forward pass start.
+
+    For sequential CFG (cfg_parallel_size=1) with PP, two full forward chains are
+    executed — one for the positive pass and one for the negative pass — so that each
+    PP stage operates on the correct encoder_hidden_states.
+    """
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+
+        if not issubclass(cls, CFGParallelMixin):
+            raise TypeError(
+                f"{cls.__name__} inherits PipelineParallelMixin but not CFGParallelMixin. "
+                "Pipeline Parallelism requires CFGParallelMixin for predict_noise(), "
+                "predict_noise_maybe_with_cfg(), scheduler_step_maybe_with_cfg(), and combine_cfg_noise(). "
+                "Add CFGParallelMixin to the base classes of your pipeline."
+            )
+        mro = cls.mro()
+        if mro.index(PipelineParallelMixin) > mro.index(CFGParallelMixin):
+            raise TypeError(
+                f"{cls.__name__} must inherit PipelineParallelMixin before CFGParallelMixin so MRO selects "
+                f"PP-aware predict/scheduler wrappers and their `super()` calls delegate to CFGParallelMixin."
+            )
+
+        init = cls.__dict__.get("__init__")
+        if callable(init):
+
+            @wraps(init)
+            def wrapped_init(self, *args: Any, **kwargs: Any) -> None:
+                init(self, *args, **kwargs)
+                vae = getattr(self, "vae", None)
+                if vae is not None and hasattr(vae, "decode"):
+                    self._wrapped_vae_decode()
+
+            cls.__init__ = wrapped_init
+
+        diffuse = cls.__dict__.get("diffuse")
+        if callable(diffuse):
+
+            @wraps(diffuse)
+            def wrapped_diffuse(self, *args: Any, **kwargs: Any) -> Any:
+                try:
+                    latents = diffuse(self, *args, **kwargs)
+                    if isinstance(latents, AsyncLatents):
+                        latents = torch.as_tensor(latents)  # avoid copying
+                    return latents
+                finally:
+                    self._sync_pp_send()
+
+            cls.diffuse = wrapped_diffuse
+
+    def _wrapped_vae_decode(self) -> None:
+        vae, orig_decode = self.vae, self.vae.decode
+
+        @wraps(orig_decode)
+        def wrapped_decode(z: torch.Tensor, *args: Any, **kwargs: Any):
+            if hasattr(vae, "is_distributed_enabled") and vae.is_distributed_enabled():
+                # Middle ranks (world size 3 or more) hold stale latents after the denoising loop.
+                # Broadcast from rank 0 so every rank splits identical tiles.
+                if get_pipeline_parallel_world_size() > 2:
+                    z = get_pp_group().broadcast(z, src=0)
+                return orig_decode(z, *args, **kwargs)
+            elif is_pipeline_first_stage():
+                return orig_decode(z, *args, **kwargs)
+            return (None,)  # decoder returns a tuple
+
+        self.vae.decode = wrapped_decode
+
+    @property
+    def _pp_send_work(self) -> list[torch.distributed.Work]:
+        if not hasattr(self, "_pp_send_work_list"):
+            self._pp_send_work_list: list[torch.distributed.Work] = []
+        return self._pp_send_work_list
+
+    @_pp_send_work.setter
+    def _pp_send_work(self, work: list[torch.distributed.Work]) -> None:
+        self._pp_send_work_list = work
+
+    def _sync_pp_send(self) -> None:
+        """
+        Wait on all pending non-blocking PP sends.
+
+        Must be called after the denoising loop so that the isend handles
+        from the last iteration are completed before any subsequent
+        collective (e.g. VAE decode broadcast) or tensor reuse.
+        """
+        if self._pp_send_work:
+            for handle in self._pp_send_work:
+                handle.wait()
+            self._pp_send_work = []
+
+    def predict_noise_maybe_with_cfg(
+        self,
+        do_true_cfg: bool,
+        true_cfg_scale: float,
+        positive_kwargs: dict[str, Any],
+        negative_kwargs: dict[str, Any] | None,
+        cfg_normalize: bool = True,
+        output_slice: int | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...] | None:
+        """
+        Drop-in replacement for predict_noise_maybe_with_cfg that also handles PP.
+
+        Supports three modes:
+          - PP only, sequential CFG: both branches (cond and uncond) run through this PP pipeline.
+            This doubles communication volume per denoising step compared to PP + CFG-parallel.
+          - PP + CFG-parallel: each PP pipeline carries one branch. The last PP
+            rank all-gathers across the CFG group and combines, mirroring
+            CFGParallelMixin.predict_noise_maybe_with_cfg exactly.
+          - PP only, no CFG: cond branch only.
+
+        Returns:
+            noise_pred on the last PP rank (all CFG ranks when CFG-parallel is active).
+            None on all other ranks.
+        """
+        if get_pipeline_parallel_world_size() == 1:
+            return super().predict_noise_maybe_with_cfg(
+                do_true_cfg, true_cfg_scale, positive_kwargs, negative_kwargs, cfg_normalize, output_slice
+            )
+
+        self._sync_pp_send()
+
+        pp_group = get_pp_group()
+
+        cfg_parallel_ready = do_true_cfg and get_classifier_free_guidance_world_size() > 1
+        if cfg_parallel_ready:
+            # Each PP pipeline carries exactly one CFG branch determined by cfg_rank.
+            all_kwargs = [positive_kwargs if get_classifier_free_guidance_rank() == 0 else negative_kwargs]
+        else:
+            # Sequential CFG (or no CFG): this PP pipeline handles all branches.
+            all_kwargs = [positive_kwargs] + ([negative_kwargs] if do_true_cfg else [])
+
+        # Non-first ranks receive intermediate tensors asynchronously
+        n = len(all_kwargs)
+        its: list[AsyncIntermediateTensors | None] = [None] * n
+        if not pp_group.is_first_rank:
+            for i in range(n):
+                its[i] = AsyncIntermediateTensors(*pp_group.irecv_tensor_dict())
+
+        if not pp_group.is_last_rank:
+            # First / middle rank: run partial forwards and propagate ITs downstream.
+            for kwargs, it in zip(all_kwargs, its):
+                result = self.predict_noise(**kwargs, intermediate_tensors=it)
+                self._pp_send_work.extend(pp_group.isend_tensor_dict(result.tensors))
+            return None
+
+        # Last rank: run full forward
+        noise_preds = [self.predict_noise(**kwargs, intermediate_tensors=it) for kwargs, it in zip(all_kwargs, its)]
+
+        if cfg_parallel_ready:
+            # All-gather the single-branch prediction across the CFG group and combine
+            # on all CFG ranks so every last PP rank has an identical noise_pred.
+            local_pred = noise_preds[0]
+            if output_slice is not None:
+                local_pred = local_pred[:, :output_slice]
+            gathered = get_cfg_group().all_gather(local_pred, separate_tensors=True)
+            return self.combine_cfg_noise(gathered[0], gathered[1], true_cfg_scale, cfg_normalize)
+
+        # Sequential CFG or no-CFG path.
+        if do_true_cfg:
+            pos, neg = noise_preds[0], noise_preds[1]
+            if output_slice is not None:
+                pos = pos[:, :output_slice]
+                neg = neg[:, :output_slice]
+            return self.combine_cfg_noise(pos, neg, true_cfg_scale, cfg_normalize)
+        pred = noise_preds[0]
+        if output_slice is not None:
+            pred = pred[:, :output_slice]
+        return pred
+
+    def scheduler_step_maybe_with_cfg(
+        self,
+        noise_pred: torch.Tensor | tuple[torch.Tensor, ...] | None,
+        t: torch.Tensor | tuple[torch.Tensor, ...],
+        latents: torch.Tensor | tuple[torch.Tensor, ...],
+        do_true_cfg: bool,
+        per_request_scheduler: Any | None = None,
+        generator: torch.Generator | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...] | AsyncLatents:
+        """
+        Drop-in replacement for scheduler_step_maybe_with_cfg that also handles PP.
+
+        Only the last rank runs the scheduler (it already has noise_pred); the result
+        is sent to rank 0 which needs it for the next forward pass.
+
+        Returns a ``AsyncLatents`` on rank 0 that transparently defers
+        ``handle.wait()`` until the tensor is actually consumed (via attribute
+        access or a torch operation), keeping the rank non-blocking after the
+        ``irecv`` is posted.
+        """
+        if get_pipeline_parallel_world_size() == 1:
+            return super().scheduler_step_maybe_with_cfg(
+                noise_pred, t, latents, do_true_cfg, per_request_scheduler, generator
+            )
+
+        pp_group = get_pp_group()
+        if pp_group.is_last_rank:
+            latents = super().scheduler_step_maybe_with_cfg(
+                noise_pred, t, latents, do_true_cfg, per_request_scheduler, generator
+            )
+            self._pp_send_work = pp_group.isend_tensor_dict({"latents": latents}, dst=0)
+        elif pp_group.is_first_rank:
+            latents = AsyncLatents(*pp_group.irecv_tensor_dict(src=pp_group.world_size - 1))
+        return latents
diff --git a/vllm_omni/diffusion/models/wan2_2/__init__.py b/vllm_omni/diffusion/models/wan2_2/__init__.py
index 2159a1184..4059c4fb5 100644
--- a/vllm_omni/diffusion/models/wan2_2/__init__.py
+++ b/vllm_omni/diffusion/models/wan2_2/__init__.py
@@ -22,11 +22,6 @@ from .pipeline_wan2_2_s2v import (
     get_wan22_s2v_post_process_func,
     get_wan22_s2v_pre_process_func,
 )
-from .pipeline_wan2_2_ti2v import (
-    Wan22TI2VPipeline,
-    get_wan22_ti2v_post_process_func,
-    get_wan22_ti2v_pre_process_func,
-)
 from .pipeline_wan2_2_vace import (
     Wan22VACEPipeline,
     get_wan22_vace_post_process_func,
@@ -50,9 +45,6 @@ __all__ = [
     "Wan22S2VPipeline",
     "get_wan22_s2v_post_process_func",
     "get_wan22_s2v_pre_process_func",
-    "Wan22TI2VPipeline",
-    "get_wan22_ti2v_post_process_func",
-    "get_wan22_ti2v_pre_process_func",
     "Wan22VACEPipeline",
     "get_wan22_vace_post_process_func",
     "get_wan22_vace_pre_process_func",
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
index e3286cf4f..b4438c795 100644
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
+++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
@@ -17,10 +17,12 @@ from torch import nn
 from transformers import AutoTokenizer, UMT5EncoderModel
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 from vllm.model_executor.models.utils import AutoWeightsLoader
+from vllm.sequence import IntermediateTensors
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -297,7 +299,9 @@ def get_wan22_pre_process_func(
     return pre_process_func
 
 
-class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin):
+class Wan22Pipeline(
+    nn.Module, PipelineParallelMixin, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin
+):
     def __init__(
         self,
         *,
@@ -460,7 +464,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
         attention_kwargs: dict[str, Any],
         latent_condition: torch.Tensor | None = None,
         first_frame_mask: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | AsyncLatents:
         if attention_kwargs is None:
             attention_kwargs = {}
         with self.progress_bar(total=len(timesteps)) as pbar:
@@ -858,7 +862,11 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
             output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
         )
 
-    def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor:
+    def predict_noise(
+        self,
+        current_model: nn.Module | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor | IntermediateTensors:
         """
         Forward pass through transformer to predict noise.
 
@@ -867,11 +875,12 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
             **kwargs: Arguments to pass to the transformer
 
         Returns:
-            Predicted noise tensor
+            Predicted noise tensor or IntermediateTensors on non-last PP stages.
         """
         if current_model is None:
             current_model = self.transformer
-        return current_model(**kwargs)[0]
+        result = current_model(**kwargs)
+        return result if isinstance(result, IntermediateTensors) else result[0]
 
     def encode_prompt(
         self,
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
index 0651356cc..42c4eff6a 100644
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
+++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
@@ -17,10 +17,12 @@ from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, UMT5EncoderModel
 from vllm.model_executor.models.utils import AutoWeightsLoader
+from vllm.sequence import IntermediateTensors
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -135,7 +137,12 @@ def get_wan22_i2v_pre_process_func(
 
 
 class Wan22I2VPipeline(
-    nn.Module, SupportImageInput, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin
+    nn.Module,
+    SupportImageInput,
+    PipelineParallelMixin,
+    CFGParallelMixin,
+    ProgressBarMixin,
+    DiffusionPipelineProfilerMixin,
 ):
     """
     Wan2.2 Image-to-Video Pipeline.
@@ -283,7 +290,7 @@ class Wan22I2VPipeline(
         attention_kwargs: dict[str, Any],
         condition: torch.Tensor,
         first_frame_mask: torch.Tensor,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | AsyncLatents:
         if attention_kwargs is None:
             attention_kwargs = {}
         with self.progress_bar(total=len(timesteps)) as pbar:
@@ -348,7 +355,6 @@ class Wan22I2VPipeline(
 
                 # Compute the previous noisy sample x_t -> x_t-1 with automatic CFG sync
                 latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
-
                 pbar.update()
 
         return latents
@@ -667,7 +673,11 @@ class Wan22I2VPipeline(
             output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
         )
 
-    def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor:
+    def predict_noise(
+        self,
+        current_model: nn.Module | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor | IntermediateTensors:
         """
         Forward pass through transformer to predict noise.
 
@@ -676,11 +686,12 @@ class Wan22I2VPipeline(
             **kwargs: Arguments to pass to the transformer
 
         Returns:
-            Predicted noise tensor
+            Predicted noise tensor or IntermediateTensors on non-last PP stages.
         """
         if current_model is None:
             current_model = self.transformer
-        return current_model(**kwargs)[0]
+        result = current_model(**kwargs)
+        return result if isinstance(result, IntermediateTensors) else result[0]
 
     def encode_prompt(
         self,
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py
deleted file mode 100644
index 360eb83b9..000000000
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py
+++ /dev/null
@@ -1,731 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Wan2.2 TI2V (Text-Image-to-Video) Pipeline.
-
-This pipeline supports the unified TI2V-5B model that can generate videos from:
-- Text only (T2V mode)
-- Text + Image (I2V mode)
-
-The key difference from the MoE-based I2V pipeline is:
-- Single transformer (not MoE with two transformers)
-- Uses expand_timesteps mode for image conditioning
-- No CLIP image encoder - only VAE encoding for image condition
-"""
-
-from __future__ import annotations
-
-import logging
-import os
-from collections.abc import Iterable
-from typing import Any, cast
-
-import numpy as np
-import PIL.Image
-import torch
-from diffusers.utils.torch_utils import randn_tensor
-from torch import nn
-from transformers import AutoTokenizer, UMT5EncoderModel
-from vllm.model_executor.models.utils import AutoWeightsLoader
-
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
-from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
-from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
-from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
-from vllm_omni.diffusion.models.interface import SupportImageInput
-from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
-from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
-    build_wan_scheduler,
-    create_transformer_from_config,
-    load_transformer_config,
-    resolve_wan_flow_shift,
-    resolve_wan_sample_solver,
-    retrieve_latents,
-)
-from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
-from vllm_omni.diffusion.postprocess import interpolate_video_tensor
-from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.inputs.data import OmniTextPrompt
-from vllm_omni.platforms import current_omni_platform
-
-logger = logging.getLogger(__name__)
-
-
-def get_wan22_ti2v_post_process_func(
-    od_config: OmniDiffusionConfig,
-):
-    from diffusers.video_processor import VideoProcessor
-
-    video_processor = VideoProcessor(vae_scale_factor=8)
-
-    def post_process_func(
-        video: torch.Tensor,
-        output_type: str = "np",
-        sampling_params=None,
-    ):
-        if output_type == "latent":
-            return video
-        custom_output = {}
-        if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False):
-            video, multiplier = interpolate_video_tensor(
-                video,
-                exp=sampling_params.frame_interpolation_exp,
-                scale=sampling_params.frame_interpolation_scale,
-                model_path=sampling_params.frame_interpolation_model_path,
-            )
-            custom_output["video_fps_multiplier"] = multiplier
-        return {
-            "video": video_processor.postprocess_video(video, output_type=output_type),
-            "custom_output": custom_output,
-        }
-
-    return post_process_func
-
-
-def get_wan22_ti2v_pre_process_func(
-    od_config: OmniDiffusionConfig,
-):
-    """Pre-process function for TI2V: optionally load and resize input image."""
-
-    def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
-        for i, prompt in enumerate(request.prompts):
-            multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None
-            raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
-            if isinstance(prompt, str):
-                prompt = OmniTextPrompt(prompt=prompt)
-            if "additional_information" not in prompt:
-                prompt["additional_information"] = {}
-
-            if raw_image is None:
-                raise ValueError(
-                    """No image is provided. This model requires an image to run.""",
-                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""",
-                )
-            if not isinstance(raw_image, (str, PIL.Image.Image)):
-                raise TypeError(
-                    f"""Unsupported image format {raw_image.__class__}.""",
-                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""",
-                )
-            image = PIL.Image.open(raw_image).convert("RGB") if isinstance(raw_image, str) else raw_image
-
-            # Calculate dimensions based on aspect ratio if not provided
-            if request.sampling_params.height is None or request.sampling_params.width is None:
-                # Default max area for 720P (TI2V-5B default)
-                max_area = 720 * 1280
-                aspect_ratio = image.height / image.width
-
-                # Calculate dimensions maintaining aspect ratio
-                mod_value = 16  # Must be divisible by 16
-                height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
-                width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
-
-                if request.sampling_params.height is None:
-                    request.sampling_params.height = height
-                if request.sampling_params.width is None:
-                    request.sampling_params.width = width
-
-            # Resize image to target dimensions
-            image = image.resize(
-                (request.sampling_params.width, request.sampling_params.height),  # type: ignore # Above has ensured that width & height are not None
-                PIL.Image.Resampling.LANCZOS,
-            )
-            prompt["multi_modal_data"]["image"] = image  # type: ignore # key existence already checked above
-
-            request.prompts[i] = prompt
-        return request
-
-    return pre_process_func
-
-
-class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressBarMixin):
-    """
-    Wan2.2 Text-Image-to-Video (TI2V) Pipeline.
-
-    This is a unified pipeline that supports both:
-    - Text-to-Video (T2V): when no image is provided
-    - Image-to-Video (I2V): when an image is provided
-
-    Uses expand_timesteps mode for I2V conditioning where the first frame
-    is conditioned on the input image latent.
-    """
-
-    def __init__(
-        self,
-        *,
-        od_config: OmniDiffusionConfig,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.od_config = od_config
-
-        self.device = get_local_device()
-        dtype = getattr(od_config, "dtype", torch.bfloat16)
-
-        model = od_config.model
-        local_files_only = os.path.exists(model)
-
-        # Set up weights sources for single transformer
-        self.weights_sources = [
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path=od_config.model,
-                subfolder="transformer",
-                revision=None,
-                prefix="transformer.",
-                fall_back_to_pt=True,
-            ),
-        ]
-
-        # See ``hub_prefetch.py`` for the transformers v5 subfolder race.
-        prefetch_subfolders(
-            model,
-            ["tokenizer", "text_encoder", "vae"],
-            local_files_only=local_files_only,
-        )
-
-        # Text encoder
-        self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.text_encoder = UMT5EncoderModel.from_pretrained(
-            model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
-        ).to(self.device)
-
-        # VAE
-        self.vae = OmniAutoencoderKLWan.from_pretrained(
-            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
-        ).to(self.device)
-
-        # Single transformer (TI2V uses dense 5B model, not MoE)
-        # Load config from model to get correct dimensions
-        transformer_config = load_transformer_config(model, "transformer", local_files_only)
-        self.transformer = self._create_transformer(transformer_config)
-
-        self._sample_solver = "unipc"
-        self._flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 5.0
-        self.scheduler = build_wan_scheduler(self._sample_solver, self._flow_shift)
-
-        # VAE scale factors
-        self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if hasattr(self.vae, "config") else 4
-        self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if hasattr(self.vae, "config") else 8
-
-        # TI2V always uses expand_timesteps mode
-        self.expand_timesteps = True
-
-        self._guidance_scale = None
-        self._num_timesteps = None
-        self._current_timestep = None
-
-    def _create_transformer(self, config: dict) -> WanTransformer3DModel:
-        """Create a transformer from a config dict. Respects od_config.quantization_config."""
-        quant_config = getattr(self.od_config, "quantization_config", None)
-        return create_transformer_from_config(config, quant_config=quant_config)
-
-    @property
-    def guidance_scale(self):
-        return self._guidance_scale
-
-    @property
-    def do_classifier_free_guidance(self):
-        return self._guidance_scale is not None and self._guidance_scale > 1.0
-
-    @property
-    def num_timesteps(self):
-        return self._num_timesteps
-
-    @property
-    def current_timestep(self):
-        return self._current_timestep
-
-    def diffuse(
-        self,
-        latents: torch.Tensor,
-        timesteps: torch.Tensor,
-        prompt_embeds: torch.Tensor,
-        negative_prompt_embeds: torch.Tensor | None,
-        guidance_scale: float,
-        dtype: torch.dtype,
-        attention_kwargs: dict[str, Any],
-        num_latent_frames: int,
-        latent_height: int,
-        latent_width: int,
-        latent_condition: torch.Tensor | None = None,
-        first_frame_mask: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        if attention_kwargs is None:
-            attention_kwargs = {}
-        with self.progress_bar(total=len(timesteps)) as pbar:
-            for step_idx, t in enumerate(timesteps):
-                self._current_timestep = t
-                set_forward_context_denoise_step_idx(step_idx)
-
-                # Prepare latent input
-                if latent_condition is not None:
-                    # I2V mode: blend condition with latents using mask
-                    latent_model_input = (1 - first_frame_mask) * latent_condition + first_frame_mask * latents
-                    latent_model_input = latent_model_input.to(dtype)
-
-                    # Expand timesteps for each patch (TI2V style)
-                    temp_ts = (first_frame_mask[0][0][:, ::2, ::2] * t).flatten()
-                    timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1)
-                else:
-                    # T2V mode: use latents directly
-                    latent_model_input = latents.to(dtype)
-
-                    # Expand timesteps for TI2V model architecture
-                    mask = torch.ones(1, 1, num_latent_frames, latent_height, latent_width, device=latents.device)
-                    temp_ts = (mask[0][0][:, ::2, ::2] * t).flatten()
-                    timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1)
-
-                do_true_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None
-                positive_kwargs = {
-                    "hidden_states": latent_model_input,
-                    "timestep": timestep,
-                    "encoder_hidden_states": prompt_embeds,
-                    "attention_kwargs": attention_kwargs,
-                    "return_dict": False,
-                    "current_model": self.transformer,
-                }
-                if do_true_cfg:
-                    negative_kwargs = {
-                        "hidden_states": latent_model_input,
-                        "timestep": timestep,
-                        "encoder_hidden_states": negative_prompt_embeds,
-                        "attention_kwargs": attention_kwargs,
-                        "return_dict": False,
-                        "current_model": self.transformer,
-                    }
-                else:
-                    negative_kwargs = None
-
-                noise_pred = self.predict_noise_maybe_with_cfg(
-                    do_true_cfg=do_true_cfg,
-                    true_cfg_scale=guidance_scale,
-                    positive_kwargs=positive_kwargs,
-                    negative_kwargs=negative_kwargs,
-                    cfg_normalize=False,
-                )
-
-                latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
-                pbar.update()
-
-        return latents
-
-    def forward(
-        self,
-        req: OmniDiffusionRequest,
-        prompt: str | None = None,
-        negative_prompt: str | None = None,
-        image: PIL.Image.Image | torch.Tensor | None = None,
-        height: int = 704,
-        width: int = 1280,
-        num_inference_steps: int = 40,
-        guidance_scale: float = 5.0,
-        frame_num: int = 81,
-        output_type: str | None = "np",
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        attention_kwargs: dict | None = None,
-        **kwargs,
-    ) -> DiffusionOutput:
-        # Get parameters from request or arguments
-        if len(req.prompts) > 1:
-            raise ValueError(
-                """This model only supports a single prompt, not a batched request.""",
-                """Please pass in a single prompt object or string, or a single-item list.""",
-            )
-        if len(req.prompts) == 1:  # If req.prompt is empty, default to prompt & neg_prompt in param list
-            prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt")
-            negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt")
-        if prompt is None and prompt_embeds is None:
-            raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.")
-
-        # Get image from request (optional for TI2V)
-        if image is None:
-            multi_modal_data = (
-                req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
-            )
-            raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
-            if isinstance(raw_image, list):
-                if len(raw_image) > 1:
-                    logger.warning(
-                        """Received a list of image. Only a single image is supported by this model."""
-                        """Taking only the first image for now."""
-                    )
-                raw_image = raw_image[0]
-            if raw_image is None:
-                image = None
-            elif isinstance(raw_image, str):
-                image = PIL.Image.open(raw_image)
-            else:
-                image = cast(PIL.Image.Image | torch.Tensor, raw_image)
-
-        # Default dimensions for TI2V-5B (720P)
-        height = req.sampling_params.height or height
-        width = req.sampling_params.width or width
-        num_frames = req.sampling_params.num_frames if req.sampling_params.num_frames else frame_num
-        num_steps = req.sampling_params.num_inference_steps or num_inference_steps
-
-        # Respect per-request guidance_scale when explicitly provided.
-        if req.sampling_params.guidance_scale_provided:
-            guidance_scale = req.sampling_params.guidance_scale
-
-        self._guidance_scale = guidance_scale
-
-        # Validate inputs
-        self.check_inputs(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            image=image,
-            height=height,
-            width=width,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-        )
-
-        # Adjust num_frames to be compatible with VAE temporal scaling
-        if num_frames % self.vae_scale_factor_temporal != 1:
-            num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1
-        num_frames = max(num_frames, 1)
-
-        device = self.device
-        dtype = self.transformer.dtype
-
-        # Generator setup
-        if generator is None:
-            generator = req.sampling_params.generator
-        if generator is None and req.sampling_params.seed is not None:
-            generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
-
-        # Encode prompts
-        if prompt_embeds is None:
-            prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                do_classifier_free_guidance=guidance_scale > 1.0,
-                num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
-                max_sequence_length=req.sampling_params.max_sequence_length or 512,
-                device=device,
-                dtype=dtype,
-            )
-        else:
-            prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
-            if negative_prompt_embeds is not None:
-                negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
-
-        batch_size = prompt_embeds.shape[0]
-
-        sample_solver = resolve_wan_sample_solver(req, default=self._sample_solver)
-        flow_shift = resolve_wan_flow_shift(req, self.od_config)
-        if sample_solver != self._sample_solver or abs(flow_shift - self._flow_shift) > 1e-6:
-            self.scheduler = build_wan_scheduler(sample_solver, flow_shift)
-            self._sample_solver = sample_solver
-            self._flow_shift = flow_shift
-
-        # Timesteps
-        self.scheduler.set_timesteps(num_steps, device=device)
-        timesteps = self.scheduler.timesteps
-        self._num_timesteps = len(timesteps)
-
-        # Prepare latents
-        num_channels_latents = self.transformer.config.in_channels
-        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
-        latent_height = height // self.vae_scale_factor_spatial
-        latent_width = width // self.vae_scale_factor_spatial
-
-        # Check if we have an image (I2V mode) or not (T2V mode)
-        if image is not None:
-            # I2V mode: prepare latents with image condition
-            from diffusers.video_processor import VideoProcessor
-
-            video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
-
-            if isinstance(image, PIL.Image.Image):
-                image_tensor = video_processor.preprocess(image, height=height, width=width)
-            else:
-                image_tensor = image
-            image_tensor = image_tensor.to(device=device, dtype=torch.float32)
-
-            latents, latent_condition, first_frame_mask = self.prepare_i2v_latents(
-                image=image_tensor,
-                batch_size=batch_size,
-                num_channels_latents=num_channels_latents,
-                height=height,
-                width=width,
-                num_frames=num_frames,
-                dtype=torch.float32,
-                device=device,
-                generator=generator,
-                latents=req.sampling_params.latents,
-            )
-        else:
-            # T2V mode: prepare random latents
-            latents = self.prepare_latents(
-                batch_size=batch_size,
-                num_channels_latents=num_channels_latents,
-                height=height,
-                width=width,
-                num_frames=num_frames,
-                dtype=torch.float32,
-                device=device,
-                generator=generator,
-                latents=req.sampling_params.latents,
-            )
-            latent_condition = None
-            first_frame_mask = torch.ones(
-                1, 1, num_latent_frames, latent_height, latent_width, dtype=torch.float32, device=device
-            )
-
-        if attention_kwargs is None:
-            attention_kwargs = {}
-
-        latents = self.diffuse(
-            latents=latents,
-            timesteps=timesteps,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            guidance_scale=guidance_scale,
-            dtype=dtype,
-            attention_kwargs=attention_kwargs,
-            num_latent_frames=num_latent_frames,
-            latent_height=latent_height,
-            latent_width=latent_width,
-            latent_condition=latent_condition,
-            first_frame_mask=first_frame_mask,
-        )
-
-        # Wan2.2 is prone to out of memory errors when predicting large videos
-        # so we empty the cache here to avoid OOM before vae decoding.
-        if current_omni_platform.is_available():
-            current_omni_platform.empty_cache()
-        self._current_timestep = None
-
-        # For I2V mode, blend final latents with condition
-        if latent_condition is not None:
-            latents = (1 - first_frame_mask) * latent_condition + first_frame_mask * latents
-
-        # Decode
-        if output_type == "latent":
-            output = latents
-        else:
-            latents = latents.to(self.vae.dtype)
-            latents_mean = (
-                torch.tensor(self.vae.config.latents_mean)
-                .view(1, self.vae.config.z_dim, 1, 1, 1)
-                .to(latents.device, latents.dtype)
-            )
-            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
-                latents.device, latents.dtype
-            )
-            latents = latents / latents_std + latents_mean
-            output = self.vae.decode(latents, return_dict=False)[0]
-
-        return DiffusionOutput(output=output)
-
-    def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor:
-        """
-        Forward pass through transformer to predict noise.
-
-        Args:
-            current_model: The transformer model to use
-            **kwargs: Arguments to pass to the transformer
-
-        Returns:
-            Predicted noise tensor
-        """
-        if current_model is None:
-            current_model = self.transformer
-        return current_model(**kwargs)[0]
-
-    def encode_prompt(
-        self,
-        prompt: str | list[str],
-        negative_prompt: str | list[str] | None = None,
-        do_classifier_free_guidance: bool = True,
-        num_videos_per_prompt: int = 1,
-        max_sequence_length: int = 512,
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
-    ):
-        """Encode text prompts using T5 text encoder."""
-        device = device or self.device
-        dtype = dtype or self.text_encoder.dtype
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        prompt_clean = [self._prompt_clean(p) for p in prompt]
-        batch_size = len(prompt_clean)
-
-        text_inputs = self.tokenizer(
-            prompt_clean,
-            padding="max_length",
-            max_length=max_sequence_length,
-            truncation=True,
-            add_special_tokens=True,
-            return_attention_mask=True,
-            return_tensors="pt",
-        )
-        ids, mask = text_inputs.input_ids, text_inputs.attention_mask
-        seq_lens = mask.gt(0).sum(dim=1).long()
-
-        prompt_embeds = self.text_encoder(ids.to(device), mask.to(device)).last_hidden_state
-        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
-        prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
-        prompt_embeds = torch.stack(
-            [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0
-        )
-
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
-
-        negative_prompt_embeds = None
-        if do_classifier_free_guidance:
-            negative_prompt = negative_prompt or ""
-            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
-            neg_text_inputs = self.tokenizer(
-                [self._prompt_clean(p) for p in negative_prompt],
-                padding="max_length",
-                max_length=max_sequence_length,
-                truncation=True,
-                add_special_tokens=True,
-                return_attention_mask=True,
-                return_tensors="pt",
-            )
-            ids_neg, mask_neg = neg_text_inputs.input_ids, neg_text_inputs.attention_mask
-            seq_lens_neg = mask_neg.gt(0).sum(dim=1).long()
-            negative_prompt_embeds = self.text_encoder(ids_neg.to(device), mask_neg.to(device)).last_hidden_state
-            negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
-            negative_prompt_embeds = [u[:v] for u, v in zip(negative_prompt_embeds, seq_lens_neg)]
-            negative_prompt_embeds = torch.stack(
-                [
-                    torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))])
-                    for u in negative_prompt_embeds
-                ],
-                dim=0,
-            )
-            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1)
-            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
-
-        return prompt_embeds, negative_prompt_embeds
-
-    @staticmethod
-    def _prompt_clean(text: str) -> str:
-        return " ".join(text.strip().split())
-
-    def prepare_latents(
-        self,
-        batch_size: int,
-        num_channels_latents: int,
-        height: int,
-        width: int,
-        num_frames: int,
-        dtype: torch.dtype | None,
-        device: torch.device | None,
-        generator: torch.Generator | list[torch.Generator] | None,
-        latents: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        """Prepare random latents for T2V mode."""
-        if latents is not None:
-            return latents.to(device=device, dtype=dtype)
-
-        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
-        shape = (
-            batch_size,
-            num_channels_latents,
-            num_latent_frames,
-            int(height) // self.vae_scale_factor_spatial,
-            int(width) // self.vae_scale_factor_spatial,
-        )
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(f"Generator list length {len(generator)} does not match batch size {batch_size}.")
-        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        return latents
-
-    def prepare_i2v_latents(
-        self,
-        image: torch.Tensor,
-        batch_size: int,
-        num_channels_latents: int,
-        height: int,
-        width: int,
-        num_frames: int,
-        dtype: torch.dtype | None,
-        device: torch.device | None,
-        generator: torch.Generator | list[torch.Generator] | None,
-        latents: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Prepare latents for I2V mode with image conditioning.
-
-        Returns:
-            latents: Initial noise latents
-            latent_condition: Encoded first frame condition
-            first_frame_mask: Mask (0 for first frame, 1 for rest)
-        """
-        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
-        latent_height = height // self.vae_scale_factor_spatial
-        latent_width = width // self.vae_scale_factor_spatial
-
-        shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width)
-
-        if latents is None:
-            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        else:
-            latents = latents.to(device=device, dtype=dtype)
-
-        # Prepare first frame condition
-        image = image.unsqueeze(2)  # [batch, channels, 1, height, width]
-        image = image.to(device=device, dtype=self.vae.dtype)
-
-        # Encode through VAE
-        latent_condition = retrieve_latents(self.vae.encode(image), sample_mode="argmax")
-        latent_condition = latent_condition.repeat(batch_size, 1, 1, 1, 1)
-
-        # Normalize latents
-        latents_mean = (
-            torch.tensor(self.vae.config.latents_mean)
-            .view(1, self.vae.config.z_dim, 1, 1, 1)
-            .to(latent_condition.device, latent_condition.dtype)
-        )
-        latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
-            latent_condition.device, latent_condition.dtype
-        )
-        latent_condition = (latent_condition - latents_mean) * latents_std
-        latent_condition = latent_condition.to(dtype)
-
-        # Create mask: 0 for first frame (condition), 1 for rest (to denoise)
-        first_frame_mask = torch.ones(1, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device)
-        first_frame_mask[:, :, 0] = 0
-
-        return latents, latent_condition, first_frame_mask
-
-    def check_inputs(
-        self,
-        prompt,
-        negative_prompt,
-        image,
-        height,
-        width,
-        prompt_embeds=None,
-        negative_prompt_embeds=None,
-    ):
-        if height % 16 != 0 or width % 16 != 0:
-            raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
-
-        if prompt is not None and prompt_embeds is not None:
-            raise ValueError("Cannot forward both `prompt` and `prompt_embeds`. Please provide only one.")
-
-        if negative_prompt is not None and negative_prompt_embeds is not None:
-            raise ValueError(
-                "Cannot forward both `negative_prompt` and `negative_prompt_embeds`. Please provide only one."
-            )
-
-        if prompt is None and prompt_embeds is None:
-            raise ValueError("Provide either `prompt` or `prompt_embeds`.")
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        """Load weights using AutoWeightsLoader for vLLM integration."""
-        loader = AutoWeightsLoader(self)
-        return loader.load_weights(weights)
diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py
index b2dcec444..81889607a 100644
--- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py
+++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py
@@ -21,9 +21,21 @@ from vllm.model_executor.layers.conv import Conv3dLayer
 from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.utils import (
+    PPMissingLayer,
+    is_pp_missing_parameter,
+    make_empty_intermediate_tensors_factory,
+    make_layers,
+)
+from vllm.sequence import IntermediateTensors
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_pipeline_parallel_world_size,
+    is_pipeline_first_stage,
+    is_pipeline_last_stage,
+)
 from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
     SequenceParallelOutput,
@@ -871,14 +883,23 @@ class WanTransformer3DModel(nn.Module):
         inner_dim = num_attention_heads * attention_head_dim
         out_channels = out_channels or in_channels
 
+        if get_pipeline_parallel_world_size() == 0 and not is_pipeline_first_stage():
+            raise RuntimeError(
+                "`initialize_model_parallel()` must be called before constructing `WanTransformer3DModel`"
+            )
+
         # 1. Patch & position embedding
         self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len)
-        self.patch_embedding = Conv3dLayer(
-            in_channels=in_channels,
-            out_channels=inner_dim,
-            kernel_size=patch_size,
-            stride=patch_size,
-        )
+        # Patch embedding only on first PP stage; other stages receive hidden_states via P2P
+        if is_pipeline_first_stage():
+            self.patch_embedding = Conv3dLayer(
+                in_channels=in_channels,
+                out_channels=inner_dim,
+                kernel_size=patch_size,
+                stride=patch_size,
+            )
+        else:
+            self.patch_embedding = PPMissingLayer()
 
         # 2. Condition embeddings
         self.condition_embedder = WanTimeTextImageEmbedding(
@@ -890,30 +911,39 @@ class WanTransformer3DModel(nn.Module):
             pos_embed_seq_len=pos_embed_seq_len,
         )
 
-        # 3. Transformer blocks
-        self.blocks = nn.ModuleList(
-            [
-                WanTransformerBlock(
-                    inner_dim,
-                    ffn_dim,
-                    num_attention_heads,
-                    eps,
-                    added_kv_proj_dim,
-                    cross_attn_norm,
-                    quant_config=quant_config,
-                    prefix=f"blocks.{layer_idx}",
-                )
-                for layer_idx in range(num_layers)
-            ]
+        # 3. Transformer blocks — partitioned across PP stages via vLLM's `make_layers`.
+        # It computes the [start_layer, end_layer) slice for this rank and fills the remaining slots
+        # with PPMissingLayer so that weight names stay globally consistent.
+        self.start_layer, self.end_layer, self.blocks = make_layers(
+            num_layers,
+            lambda prefix: WanTransformerBlock(
+                inner_dim,
+                ffn_dim,
+                num_attention_heads,
+                eps,
+                added_kv_proj_dim,
+                cross_attn_norm,
+                quant_config=quant_config,
+                prefix=prefix,
+            ),
+            prefix="blocks",
         )
 
-        # 4. Output norm & projection
-        self.norm_out = AdaLayerNorm(inner_dim, elementwise_affine=False, eps=eps)
-        self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))
+        # 4. Output norm & projection — only on the last PP stage
+        if is_pipeline_last_stage():
+            self.norm_out = AdaLayerNorm(inner_dim, elementwise_affine=False, eps=eps)
+            self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))
+        else:
+            self.norm_out, self.proj_out = PPMissingLayer(), PPMissingLayer()
 
         # SP helper modules
         self.timestep_proj_prepare = TimestepProjPrepare()
-        self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim)
+        if is_pipeline_last_stage():
+            self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim)
+        else:
+            self.output_scale_shift_prepare = PPMissingLayer()
+
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(["hidden_states"], inner_dim)
 
         # ROPE helper
         self._cached_rope_emb = None
@@ -930,9 +960,10 @@ class WanTransformer3DModel(nn.Module):
         timestep: torch.LongTensor,
         encoder_hidden_states: torch.Tensor,
         encoder_hidden_states_image: torch.Tensor | None = None,
+        intermediate_tensors: IntermediateTensors | None = None,
         return_dict: bool = True,
         attention_kwargs: dict[str, Any] | None = None,
-    ) -> torch.Tensor | Transformer2DModelOutput:
+    ) -> torch.Tensor | Transformer2DModelOutput | IntermediateTensors:
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         p_t, p_h, p_w = self.config.patch_size
         post_patch_num_frames = num_frames // p_t
@@ -949,18 +980,25 @@ class WanTransformer3DModel(nn.Module):
             self._hidden_states_shape = hidden_states.shape
             self._cached_rope_emb = rotary_emb
 
-        # Patch embedding and flatten to sequence
-        # (hidden_states is sharded at blocks.0 input by _sp_plan)
-        hidden_states = self.patch_embedding(hidden_states)
-        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+        if is_pipeline_first_stage():
+            # Patch embedding and flatten to sequence
+            # (hidden_states is sharded at blocks.0 input by _sp_plan)
+            hidden_states = self.patch_embedding(hidden_states)
+            hidden_states = hidden_states.flatten(2).transpose(1, 2)
+        else:
+            if intermediate_tensors is None:
+                raise RuntimeError("intermediate_tensors must be provided for non-first PP stages")
+            hidden_states = intermediate_tensors["hidden_states"]
 
-        # Handle timestep shape
+        # Handle timestep shape (2-D for TI2V, 1-D for T2V)
         if timestep.ndim == 2:
             ts_seq_len = timestep.shape[1]
             timestep = timestep.flatten()
         else:
             ts_seq_len = None
 
+        # Compute conditioning on all PP stages.
+        # Each stage needs temb/timestep_proj for scale-shift modulation in its local blocks.
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
             timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len
         )
@@ -996,14 +1034,13 @@ class WanTransformer3DModel(nn.Module):
             hidden_states_mask = None
 
         # Transformer blocks
-        for block in self.blocks:
-            hidden_states = block(
-                hidden_states,
-                encoder_hidden_states,
-                timestep_proj,
-                rotary_emb,
-                hidden_states_mask,
-            )
+        for block in self.blocks[self.start_layer : self.end_layer]:
+            hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, hidden_states_mask)
+
+        if not is_pipeline_last_stage():
+            # Non-last PP stage: hand the token sequence to the caller via IntermediateTensors.
+            # predict_noise will broadcast it to the next stage before calling that stage's forward.
+            return IntermediateTensors({"hidden_states": hidden_states})
 
         # Output norm, projection & unpatchify
         shift, scale = self.output_scale_shift_prepare(temb)
@@ -1074,7 +1111,8 @@ class WanTransformer3DModel(nn.Module):
                 if weight_name not in original_name:
                     continue
                 lookup_name = original_name.replace(weight_name, param_name)
-                if lookup_name not in params_dict:
+                # Skip weights that belong to PP stages other than this one
+                if is_pp_missing_parameter(lookup_name, self) or lookup_name not in params_dict:
                     break
                 param = params_dict[lookup_name]
                 weight_loader = param.weight_loader
@@ -1098,6 +1136,10 @@ class WanTransformer3DModel(nn.Module):
                     if modulation_alias in params_dict:
                         lookup_name = modulation_alias
 
+                # Skip weights that belong to PP stages other than this one
+                if is_pp_missing_parameter(lookup_name, self):
+                    continue
+
                 if lookup_name not in params_dict:
                     logger.warning(f"Skipping weight {original_name} -> {lookup_name}")
                     continue
diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py
index 687aa3bda..8a850f565 100644
--- a/vllm_omni/engine/async_omni_engine.py
+++ b/vllm_omni/engine/async_omni_engine.py
@@ -1632,6 +1632,7 @@ class AsyncOmniEngine:
             data_parallel_size = normalized_kwargs.get("data_parallel_size") or 1
             tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1
             cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1
+            pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
             vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1
             enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False
             use_hsdp = normalized_kwargs.get("use_hsdp", False)