Merge branch 'main' into lance-model-support

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -136,7 +136,9 @@ class FlashAttentionImpl(AttentionImpl):
         )
 
         batch_size, q_len = query.size()[:2]
-        cu_seqlens = torch.arange(0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=query.device)
+        k_len = key.size(1)
+        cu_seqlens_q = torch.arange(0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=query.device)
+        cu_seqlens_k = torch.arange(0, (batch_size + 1) * k_len, step=k_len, dtype=torch.int32, device=query.device)
         # b s ... -> (b s) ...
         query = query.flatten(0, 1)
         key = key.flatten(0, 1)
@@ -146,10 +148,10 @@ class FlashAttentionImpl(AttentionImpl):
             q=query,
             k=key,
             v=value,
-            cu_seqlens_q=cu_seqlens,
-            cu_seqlens_k=cu_seqlens,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
             max_seqlen_q=q_len,
-            max_seqlen_k=q_len,
+            max_seqlen_k=k_len,
             causal=self.causal,
             softmax_scale=self.softmax_scale,
         )

vllm_omni/diffusion/attention/layer.py

@@ -171,9 +171,13 @@ class Attention(nn.Module):
     def _init_kv_cache_quantization(self, config) -> None:
         if config is None:
             return
-        dtype = config.diffusion_kv_cache_dtype
+        dtype = getattr(config, "diffusion_kv_cache_dtype", None)
+        if dtype == "auto":
+            dtype = None
+        parallel_config = getattr(config, "parallel_config", None)
+        ring_degree = getattr(parallel_config, "ring_degree", 1)
         if dtype:
-            if config.parallel_config.ring_degree > 1:
+            if ring_degree > 1:
                 raise ValueError(
                     "KV quantization is not compatible with ring attention "
                     "(ring_degree > 1). Ring kernels do not propagate quantization descale "
@@ -190,8 +194,8 @@ class Attention(nn.Module):
                 )
                 dtype = None
         self._kv_cache_dtype = dtype
-        self._kv_cache_skip_steps = config.diffusion_kv_cache_skip_step_indices
-        self._kv_cache_skip_layers = config.diffusion_kv_cache_skip_layer_indices
+        self._kv_cache_skip_steps = getattr(config, "diffusion_kv_cache_skip_step_indices", None)
+        self._kv_cache_skip_layers = getattr(config, "diffusion_kv_cache_skip_layer_indices", None)
 
     def _should_apply_kv_cache_quant(self) -> bool:
         skip_steps = self._kv_cache_skip_steps

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,

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_hunyuan.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+from vllm_omni.diffusion.models.hunyuan_image3.autoencoder import AutoencoderKLConv3D
+
+logger = init_logger(__name__)
+
+
+class DistributedAutoencoderKLHunyuan(AutoencoderKLConv3D, DistributedVaeMixin):
+    @classmethod
+    def from_config(cls, config: Any, **kwargs: Any):
+        model = super().from_config(config, **kwargs)
+        model.init_distributed()
+        return model
+
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+
+    @property
+    def use_tiling(self) -> bool:
+        return self.use_spatial_tiling
+
+    @use_tiling.setter
+    def use_tiling(self, use_tiling: bool) -> None:
+        self.use_spatial_tiling = use_tiling
+
+    def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        _, _, _, height, width = z.shape
+        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
+        row_limit = int(self.tile_sample_min_size - blend_extent)
+
+        tiletask_list = []
+        for i in range(0, height, overlap_size):
+            for j in range(0, width, overlap_size):
+                tile = z[:, :, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
+                tiletask_list.append(
+                    TileTask(
+                        len(tiletask_list),
+                        (i // overlap_size, j // overlap_size),
+                        tile,
+                        workload=tile.shape[3] * tile.shape[4],
+                    )
+                )
+
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec={"blend_extent": blend_extent, "row_limit": row_limit},
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def tile_exec(self, task: TileTask) -> torch.Tensor:
+        return self.decoder(task.tensor)
+
+    def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
+        grid_h, grid_w = grid_spec.grid_shape
+        result_rows = []
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_extent"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_extent"])
+                result_row.append(tile[:, :, :, : grid_spec.tile_spec["row_limit"], : grid_spec.tile_spec["row_limit"]])
+            result_rows.append(torch.cat(result_row, dim=-1))
+        return torch.cat(result_rows, dim=-2)
+
+    def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        _, _, _, height, width = x.shape
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = int(self.tile_latent_min_size - blend_extent)
+
+        tiletask_list = []
+        for i in range(0, height, overlap_size):
+            for j in range(0, width, overlap_size):
+                tile = x[:, :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
+                tiletask_list.append(
+                    TileTask(
+                        len(tiletask_list),
+                        (i // overlap_size, j // overlap_size),
+                        tile,
+                        workload=tile.shape[3] * tile.shape[4],
+                    )
+                )
+
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec={"blend_extent": blend_extent, "row_limit": row_limit},
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def encode_tile_exec(self, task: TileTask) -> torch.Tensor:
+        return self.encoder(task.tensor)
+
+    def encode_tile_merge(
+        self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec
+    ) -> torch.Tensor:
+        return self.tile_merge(coord_tensor_map, grid_spec)
+
+    def spatial_tiled_encode(self, x: torch.Tensor):
+        if not self.is_distributed_enabled():
+            return super().spatial_tiled_encode(x)
+
+        logger.debug("Encode running with distributed executor")
+        return self.distributed_executor.execute(
+            x,
+            DistributedOperator(
+                split=self.encode_tile_split,
+                exec=self.encode_tile_exec,
+                merge=self.encode_tile_merge,
+            ),
+            broadcast_result=True,
+        )
+
+    def spatial_tiled_decode(self, z: torch.Tensor):
+        if not self.is_distributed_enabled():
+            return super().spatial_tiled_decode(z)
+
+        logger.debug("Decode running with distributed executor")
+        return self.distributed_executor.execute(
+            z,
+            DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
+            broadcast_result=True,
+        )

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,

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],

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
     )

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

vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py

@@ -701,6 +701,10 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
             raise ValueError(
                 "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
             )
+        elif isinstance(prompt, str) and not prompt.strip():
+            raise ValueError("`prompt` cannot be empty or whitespace-only.")
+        elif isinstance(prompt, list) and any(isinstance(p, str) and not p.strip() for p in prompt):
+            raise ValueError("`prompt` cannot contain empty or whitespace-only strings.")
         elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
             raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
 
@@ -861,7 +865,11 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
 
         height = req.sampling_params.height or height
         width = req.sampling_params.width or width
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        num_inference_steps = (
+            req.sampling_params.num_inference_steps
+            if req.sampling_params.num_inference_steps is not None
+            else num_inference_steps
+        )
         sigmas = req.sampling_params.sigmas or sigmas
         guidance_scale = (
             req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -72,6 +72,7 @@ from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelOutput,
 )
 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.layers.rope import RotaryEmbedding
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import HunyuanFusedMoE
 from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding
@@ -868,7 +869,15 @@ class ImageKVCacheManager:
     Manages specialized caching and updating of KV-Cache for image tokens in multimodal models.
     """
 
-    def __init__(self, num_heads: int, num_kv_heads: int, head_dim: int, scaling: float, image_token_len: int = 4097):
+    def __init__(
+        self,
+        num_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        scaling: float,
+        image_token_len: int = 4097,
+        prefix: str = "",
+    ):
         """
         Args:
             image_token_len: Number of tokens per image (including special placeholders),
@@ -892,6 +901,7 @@ class ImageKVCacheManager:
             causal=False,
             softmax_scale=self.scaling,
             num_kv_heads=self.num_kv_heads,
+            prefix=f"{prefix}.attn" if prefix else "",
         )
 
     def _cache_prompt_kv(
@@ -1650,6 +1660,7 @@ class HunYuanAttention(nn.Module):
             num_kv_heads=self.num_kv_heads,
             scaling=self.scaling,
             image_token_len=4097,
+            prefix=f"{prefix}.image_attn",
         )
         self.image_rope2d_emb = HunYuanRotary2DEmbedder(
             num_heads=self.num_heads,
@@ -3021,6 +3032,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
 
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
+                set_forward_context_denoise_step_idx(i)
                 if cfg_parallel_ready:
                     # CFG parallel: each rank forwards its own branch (no batch doubling)
                     latent_model_input = latents
@@ -3109,6 +3121,8 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                     progress_bar.update()
 
+        set_forward_context_denoise_step_idx(None)
+
         if hasattr(self.vae.config, "scaling_factor") and self.vae.config.scaling_factor:
             latents = latents / self.vae.config.scaling_factor
         if hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor:

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -30,7 +30,6 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import Siglip2VisionTransformer
 
-from .autoencoder import AutoencoderKLConv3D
 from .hunyuan_image3_tokenizer import TokenizerWrapper
 from .hunyuan_image3_transformer import (
     CausalMMOutputWithPast,
@@ -343,7 +342,13 @@ class HunyuanImage3Pipeline(
         quant_config = od_config.quantization_config
         self.model = HunyuanImage3Model(self.hf_config, quant_config=quant_config)
         self.transformer = self.model
-        self.vae = AutoencoderKLConv3D.from_config(self.hf_config.vae)
+        # Lazy import to break circular dependency:
+        # autoencoder_kl_hunyuan -> hunyuan_image3/__init__ -> pipeline_hunyuan_image3 -> autoencoder_kl_hunyuan
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan import (  # noqa: PLC0415
+            DistributedAutoencoderKLHunyuan,
+        )
+
+        self.vae = DistributedAutoencoderKLHunyuan.from_config(self.hf_config.vae)
         self.vae.use_spatial_tiling = self.od_config.vae_use_tiling
         self._pipeline = None
         self._tkwrapper = TokenizerWrapper(od_config.model)

vllm_omni/diffusion/models/ltx2/ltx2_transformer.py

@@ -17,7 +17,7 @@ import inspect
 from collections.abc import Iterable
 from dataclasses import dataclass
 from types import SimpleNamespace
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.distributed
@@ -39,6 +39,9 @@ from vllm.model_executor.layers.linear import (
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
@@ -173,7 +176,16 @@ class LTX2AdaLayerNormSingle(nn.Module):
 
 
 class ColumnParallelApproxGELU(nn.Module):
-    def __init__(self, dim_in: int, dim_out: int, *, approximate: str, bias: bool = True):
+    def __init__(
+        self,
+        dim_in: int,
+        dim_out: int,
+        *,
+        approximate: str,
+        bias: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ):
         super().__init__()
         self.proj = ColumnParallelLinear(
             dim_in,
@@ -181,6 +193,8 @@ class ColumnParallelApproxGELU(nn.Module):
             bias=bias,
             gather_output=False,
             return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.proj" if prefix else "proj",
         )
         self.approximate = approximate
 
@@ -200,6 +214,8 @@ class LTX2FeedForward(nn.Module):
         bias: bool = True,
         dropout: float = 0.0,
         final_dropout: bool = False,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
 
@@ -211,13 +227,22 @@ class LTX2FeedForward(nn.Module):
         dropout_layer: nn.Module = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
 
         layers: list[nn.Module] = [
-            ColumnParallelApproxGELU(dim, inner_dim, approximate="tanh", bias=bias),
+            ColumnParallelApproxGELU(
+                dim,
+                inner_dim,
+                approximate="tanh",
+                bias=bias,
+                quant_config=quant_config,
+                prefix=f"{prefix}.net.0" if prefix else "net.0",
+            ),
             dropout_layer,
             RowParallelLinear(
                 inner_dim,
                 dim_out,
                 input_is_parallel=True,
                 return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.net.2" if prefix else "net.2",
             ),
         ]
         if final_dropout:
@@ -515,6 +540,9 @@ class LTX2Attention(torch.nn.Module):
         rope_type: str = "interleaved",
         apply_gated_attention: bool = False,
         processor=None,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+        disable_kv_quant: bool = False,
     ):
         super().__init__()
         # LTX-2 uses "rms_norm_across_heads", LTX-2.3 uses "rms_norm" -- both
@@ -548,6 +576,8 @@ class LTX2Attention(torch.nn.Module):
                 head_size=self.head_dim,
                 total_num_heads=heads,
                 bias=bias,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_qkv" if prefix else "to_qkv",
             )
             self.query_num_heads = self.to_qkv.num_heads
             self.kv_num_heads = self.to_qkv.num_kv_heads
@@ -562,6 +592,8 @@ class LTX2Attention(torch.nn.Module):
                 bias=bias,
                 gather_output=False,
                 return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_q" if prefix else "to_q",
             )
             self.to_k = ColumnParallelLinear(
                 self.cross_attention_dim,
@@ -569,6 +601,8 @@ class LTX2Attention(torch.nn.Module):
                 bias=bias,
                 gather_output=False,
                 return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_k" if prefix else "to_k",
             )
             self.to_v = ColumnParallelLinear(
                 self.cross_attention_dim,
@@ -576,6 +610,8 @@ class LTX2Attention(torch.nn.Module):
                 bias=bias,
                 gather_output=False,
                 return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_v" if prefix else "to_v",
             )
 
         self.heads = self.query_num_heads
@@ -617,6 +653,8 @@ class LTX2Attention(torch.nn.Module):
                     bias=out_bias,
                     input_is_parallel=True,
                     return_bias=False,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.to_out.0" if prefix else "to_out.0",
                 ),
                 torch.nn.Dropout(dropout) if dropout > 0 else torch.nn.Identity(),
             ]
@@ -627,9 +665,12 @@ class LTX2Attention(torch.nn.Module):
             num_kv_heads=self.kv_num_heads,
             softmax_scale=1.0 / (dim_head**0.5),
             causal=False,
+            prefix=prefix,
+            disable_kv_quant=disable_kv_quant,
         )
 
         # LTX-2.3: per-head gated attention
+        # leave unquantized for this linear
         if apply_gated_attention:
             self.to_gate_logits = nn.Linear(query_dim, heads, bias=True)
         else:
@@ -748,6 +789,8 @@ class LTX2VideoTransformerBlock(nn.Module):
         elementwise_affine: bool = False,
         rope_type: str = "interleaved",
         perturbed_attn: bool = False,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.video_cross_attn_adaln = video_cross_attn_adaln
@@ -767,6 +810,8 @@ class LTX2VideoTransformerBlock(nn.Module):
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=video_gated_attn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn1" if prefix else "attn1",
         )
 
         self.audio_norm1 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -781,6 +826,8 @@ class LTX2VideoTransformerBlock(nn.Module):
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=audio_gated_attn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.audio_attn1" if prefix else "audio_attn1",
         )
 
         # 2. Prompt Cross-Attention
@@ -796,6 +843,9 @@ class LTX2VideoTransformerBlock(nn.Module):
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=video_gated_attn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn2" if prefix else "attn2",
+            disable_kv_quant=True,
         )
 
         self.audio_norm2 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -810,9 +860,14 @@ class LTX2VideoTransformerBlock(nn.Module):
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=audio_gated_attn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.audio_attn2" if prefix else "audio_attn2",
+            disable_kv_quant=True,
         )
 
         # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
+        # K/V here come from the other modality's latents (long sequences),
+        # so KV-cache quant remains enabled.
         self.audio_to_video_norm = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine)
         self.audio_to_video_attn = LTX2Attention(
             query_dim=dim,
@@ -825,6 +880,8 @@ class LTX2VideoTransformerBlock(nn.Module):
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=video_gated_attn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.audio_to_video_attn" if prefix else "audio_to_video_attn",
         )
 
         self.video_to_audio_norm = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -839,14 +896,26 @@ class LTX2VideoTransformerBlock(nn.Module):
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=audio_gated_attn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.video_to_audio_attn" if prefix else "video_to_audio_attn",
         )
 
         # 4. Feedforward layers
         self.norm3 = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine)
-        self.ff = LTX2FeedForward(dim, activation_fn=activation_fn)
+        self.ff = LTX2FeedForward(
+            dim,
+            activation_fn=activation_fn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.ff" if prefix else "ff",
+        )
 
         self.audio_norm3 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
-        self.audio_ff = LTX2FeedForward(audio_dim, activation_fn=activation_fn)
+        self.audio_ff = LTX2FeedForward(
+            audio_dim,
+            activation_fn=activation_fn,
+            quant_config=quant_config,
+            prefix=f"{prefix}.audio_ff" if prefix else "audio_ff",
+        )
 
         # 5. Per-Layer Modulation Parameters
         # LTX-2.3 with cross_attn_adaln uses 9 params (extra 3 for cross-attn modulation);
@@ -1365,6 +1434,9 @@ class LTX2VideoTransformer3DModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["transformer_blocks"]
     _hsdp_shard_conditions = [is_transformer_block_module]
     _sp_plan: dict[str, Any] | None = None
+    packed_modules_mapping = {
+        "to_qkv": ["to_q", "to_k", "to_v"],
+    }
 
     @staticmethod
     def _build_sp_plan(rope_type: str) -> dict[str, Any]:
@@ -1453,6 +1525,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
         cross_attn_mod: bool = False,
         audio_gated_attn: bool = False,
         audio_cross_attn_mod: bool = False,
+        quant_config: "QuantizationConfig | None" = None,
     ) -> None:
         super().__init__()
         self.perturbed_attn = perturbed_attn
@@ -1647,8 +1720,10 @@ class LTX2VideoTransformer3DModel(nn.Module):
                     elementwise_affine=norm_elementwise_affine,
                     rope_type=rope_type,
                     perturbed_attn=perturbed_attn,
+                    quant_config=quant_config,
+                    prefix=f"transformer_blocks.{layer_idx}",
                 )
-                for _ in range(num_layers)
+                for layer_idx in range(num_layers)
             ]
         )
 

vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py

@@ -9,10 +9,13 @@ import json
 import os
 from collections.abc import Iterable
 from contextlib import nullcontext
-from typing import Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import numpy as np
 import torch
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler
 from diffusers.pipelines.ltx2 import LTX2TextConnectors
 from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
@@ -67,14 +70,20 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc
     return {}
 
 
-def create_transformer_from_config(config: dict) -> LTX2VideoTransformer3DModel:
+def create_transformer_from_config(
+    config: dict,
+    quant_config: QuantizationConfig | None = None,
+) -> LTX2VideoTransformer3DModel:
     """Create LTX2VideoTransformer3DModel from config dict."""
-    if not config:
+    if not config and quant_config is None:
         return LTX2VideoTransformer3DModel()
 
     signature = inspect.signature(LTX2VideoTransformer3DModel.__init__)
     allowed_keys = set(signature.parameters.keys())
     kwargs = {k: v for k, v in config.items() if k in allowed_keys}
+    if quant_config is not None:
+        kwargs["quant_config"] = quant_config
+
     return LTX2VideoTransformer3DModel(**kwargs)
 
 
@@ -214,7 +223,8 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
         ).to(self.device)
 
         transformer_config = load_transformer_config(model, "transformer", local_files_only)
-        self.transformer = create_transformer_from_config(transformer_config)
+        quant_config = getattr(self.od_config, "quantization_config", None)
+        self.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config)
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py

@@ -174,7 +174,8 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
 
         # --- Transformer: created empty, weights loaded via AutoWeightsLoader ---
         transformer_config = load_transformer_config(model, "transformer", local_files_only)
-        self.transformer = create_transformer_from_config(transformer_config)
+        quant_config = getattr(self.od_config, "quantization_config", None)
+        self.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config)
 
         # --- Scheduler ---
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -290,7 +290,22 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
         )
         self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
             model, subfolder="text_encoder", local_files_only=local_files_only
-        ).to(self.device)
+        )
+        # Qwen2.5-VL ships a vision tower that text-to-image does not use.
+        # Drop it while the model is still on CPU, before moving to GPU, so
+        # the vision tower never consumes GPU memory. Handle both transformers
+        # layouts: newer puts visual under .model, older puts it directly on
+        # the model.
+        visual_owner = None
+        if hasattr(self.text_encoder, "model") and hasattr(self.text_encoder.model, "visual"):
+            visual_owner = self.text_encoder.model
+        elif hasattr(self.text_encoder, "visual"):
+            visual_owner = self.text_encoder
+        if visual_owner is not None:
+            del visual_owner.visual
+        else:
+            logger.warning("Qwen-Image: vision tower not found on text encoder; skipping drop")
+        self.text_encoder = self.text_encoder.to(self.device)
         self.vae = DistributedAutoencoderKLQwenImage.from_pretrained(
             model, subfolder="vae", local_files_only=local_files_only
         ).to(self.device)

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",

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
@@ -120,8 +122,7 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc
 
 
 def create_transformer_from_config(
-    config: dict,
-    quant_config: QuantizationConfig | None = None,
+    config: dict, quant_config: QuantizationConfig | None = None, prefix: str = ""
 ) -> WanTransformer3DModel:
     """Create WanTransformer3DModel from config dict."""
     kwargs: dict = {}
@@ -157,57 +158,15 @@ def create_transformer_from_config(
     if "pos_embed_seq_len" in config:
         kwargs["pos_embed_seq_len"] = config["pos_embed_seq_len"]
 
-    # Auto-detect quantization from transformer's config.json when not explicitly provided.
-    # merge_mxfp8_checkpoint.py injects quantization_config into config.json so that
-    # offline quantized checkpoints are recognized here without a CLI flag.
     if "quantization_config" in config:
-        from vllm_omni.quantization.factory import build_quant_config
+        from vllm_omni.quantization.factory import resolve_quant_config_from_disk
 
-        disk_qc = config["quantization_config"]
-        if isinstance(disk_qc, dict) and "quant_method" in disk_qc:
-            qc_method = disk_qc["quant_method"]
-            qc_kwargs = {k: v for k, v in disk_qc.items() if k != "quant_method"}
-            if quant_config is None:
-                # No CLI flag: full auto-detection.
-                quant_config = build_quant_config(qc_method, **qc_kwargs)
-                logger.info(
-                    "Auto-detected quantization from transformer config.json: method=%s kwargs=%s",
-                    qc_method,
-                    qc_kwargs,
-                )
-            elif quant_config.get_name() != qc_method:
-                # The caller supplied a quant_config of a different method than
-                # what the checkpoint was built with. Loading serialized tensors
-                # (e.g. MXFP8 weight scales) with the wrong linear method would
-                # produce corrupt output or a shape mismatch crash.  Reject early
-                # so the user gets a clear message instead of a silent failure.
-                raise ValueError(
-                    f"Checkpoint config.json declares quant_method={qc_method!r} but the "
-                    f"active quantization config is {quant_config.get_name()!r}. "
-                    "Pass a matching --quantization flag or omit it for auto-detection."
-                )
-            elif (
-                qc_kwargs.get("is_checkpoint_mxfp8_serialized", False)
-                and hasattr(quant_config, "is_checkpoint_mxfp8_serialized")
-                and not quant_config.is_checkpoint_mxfp8_serialized
-            ):
-                # Same method: CLI provided online mode but config.json marks this
-                # as a pre-quantized offline checkpoint.  Switch to offline mode so
-                # users can pass --quantization mxfp8 without knowing the
-                # online/offline distinction.
-                quant_config = build_quant_config(qc_method, **qc_kwargs)
-                logger.info(
-                    "config.json marks checkpoint as serialized; switching from online to offline MXFP8 mode.",
-                )
-        elif isinstance(disk_qc, str) and quant_config is None:
-            quant_config = build_quant_config(disk_qc)
-            logger.info(
-                "Auto-detected quantization from transformer config.json: method=%s",
-                disk_qc,
-            )
+        quant_config = resolve_quant_config_from_disk(quant_config, config["quantization_config"])
 
     if quant_config is not None:
         kwargs["quant_config"] = quant_config
+    if prefix:
+        kwargs["prefix"] = prefix
 
     return WanTransformer3DModel(**kwargs)
 
@@ -297,7 +256,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 +421,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 +819,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 +832,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,

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.
@@ -224,10 +231,25 @@ class Wan22I2VPipeline(
         # Transformers (weights loaded via load_weights)
         # Load config from model directory or HF Hub to get correct in_channels for I2V models
         transformer_config = load_transformer_config(model, "transformer", local_files_only)
-        self.transformer = self._create_transformer(transformer_config)
+        self.transformer = create_transformer_from_config(
+            transformer_config,
+            quant_config=od_config.quantization_config,
+        )
         if self.has_transformer_2:
             transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only)
-            self.transformer_2 = self._create_transformer(transformer_2_config)
+            t2_quant = transformer_2_config.get("quantization_config")
+            if isinstance(t2_quant, dict) and "quant_method" in t2_quant:
+                from vllm_omni.quantization.factory import build_quant_config
+
+                method = t2_quant["quant_method"]
+                kwargs = {k: v for k, v in t2_quant.items() if k != "quant_method"}
+                t2_quant = build_quant_config(method, **kwargs)
+            else:
+                t2_quant = None
+            self.transformer_2 = create_transformer_from_config(
+                transformer_2_config,
+                quant_config=t2_quant,
+            )
         else:
             self.transformer_2 = None
 
@@ -283,7 +305,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 +370,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 +688,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 +701,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,

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py

@@ -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)

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py

@@ -45,6 +45,7 @@ logger = init_logger(__name__)
 def create_vace_transformer_from_config(
     config: dict,
     quant_config: QuantizationConfig | None = None,
+    prefix: str = "",
 ) -> WanVACETransformer3DModel:
     """Create WanVACETransformer3DModel from config dict."""
     kwargs = {}
@@ -84,6 +85,8 @@ def create_vace_transformer_from_config(
         kwargs["vace_in_channels"] = config["vace_in_channels"]
     if quant_config is not None:
         kwargs["quant_config"] = quant_config
+    if prefix:
+        kwargs["prefix"] = prefix
 
     return WanVACETransformer3DModel(**kwargs)
 

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

vllm_omni/diffusion/sched/base_scheduler.py

@@ -22,7 +22,9 @@ from vllm_omni.diffusion.sched.interface import (
 
 logger = init_logger(__name__)
 
-_KEY_FIELD_NAMES = frozenset(f.name for f in fields(SamplingParamsKey))
+# LoRA identity is derived from `sampling.lora_request`, not a same-named field
+# on sampling params, so it must be resolved separately from the bulk lookup.
+_KEY_FIELD_NAMES = frozenset(f.name for f in fields(SamplingParamsKey)) - {"lora_int_id"}
 
 
 def get_sampling_params_key(request: OmniDiffusionRequest) -> SamplingParamsKey | None:
@@ -31,7 +33,11 @@ def get_sampling_params_key(request: OmniDiffusionRequest) -> SamplingParamsKey
         return None
 
     sampling = request.sampling_params
-    return SamplingParamsKey(**{name: getattr(sampling, name) for name in _KEY_FIELD_NAMES})
+    lora_request = getattr(sampling, "lora_request", None)
+    return SamplingParamsKey(
+        lora_int_id=lora_request.lora_int_id if lora_request is not None else None,
+        **{name: getattr(sampling, name) for name in _KEY_FIELD_NAMES},
+    )
 
 
 class _BaseScheduler(SchedulerInterface):

vllm_omni/diffusion/sched/interface.py

@@ -65,6 +65,11 @@ class SamplingParamsKey:
     true_cfg_scale: float | None = None
     cfg_normalize: bool = False
 
+    # LoRA identity. Requests with different adapters or scales must run in
+    # separate batches so the worker can activate exactly one adapter per step.
+    lora_int_id: int | None = None
+    lora_scale: float = 1.0
+
 
 @dataclass
 class DiffusionRequestState:

vllm_omni/diffusion/stage_diffusion_client.py

@@ -17,6 +17,7 @@ from threading import Thread
 from typing import TYPE_CHECKING, Any
 
 import zmq
+import zmq.asyncio
 from vllm.logger import init_logger
 from vllm.v1.engine.exceptions import EngineDeadError
 
@@ -124,6 +125,11 @@ class StageDiffusionClient(StageClientBase):
         self.engine_input_source = getattr(metadata, "engine_input_source", [])
         self._proc = proc
         self._owns_process = proc is not None
+        # Expose the ZMQ addresses on the instance so callers (e.g.
+        # ``StagePool._client_input_addr``) can identify the diffusion
+        # replica by its bound address.
+        self.request_address = request_address
+        self.response_address = response_address
 
         self._zmq_ctx = zmq.Context()
         self._request_socket = self._zmq_ctx.socket(zmq.PUSH)
@@ -131,6 +137,9 @@ class StageDiffusionClient(StageClientBase):
         self._response_socket = self._zmq_ctx.socket(zmq.PULL)
         self._response_socket.connect(response_address)
 
+        self._response_poller = zmq.asyncio.Poller()
+        self._response_poller.register(self._response_socket, zmq.POLLIN)
+
         self._encoder = OmniMsgpackEncoder()
         self._decoder = OmniMsgpackDecoder()
 
@@ -468,17 +477,27 @@ class StageDiffusionClient(StageClientBase):
         try:
             while True:
                 self._drain_responses()
-                if rpc_id in self._rpc_results:
-                    return self._rpc_results.pop(rpc_id)
+                result = self._rpc_results.pop(rpc_id, None)
+                if result is not None:
+                    return result
                 if self._engine_dead or (self._owns_process and self._proc is not None and not self._proc.is_alive()):
                     self._engine_dead = True
                     raise EngineDeadError(
                         f"StageDiffusionProc died while waiting for "
                         f"collective_rpc '{method}' (exit code {self._proc.exitcode})"
                     )
-                if deadline and time.monotonic() > deadline:
+                if deadline is not None and time.monotonic() > deadline:
                     raise TimeoutError(f"collective_rpc_async '{method}' timed out after {timeout}s")
-                await asyncio.sleep(0.01)
+                # Block (async) until data arrives on the ZMQ response
+                # socket or until the timeout expires, then loop back to
+                # drain and check.
+                if deadline is not None:
+                    poll_timeout_ms = max(int((deadline - time.monotonic()) * 1000), 0)
+                else:
+                    poll_timeout_ms = 100
+                # no exception raised on timeout (capped at 100ms so the
+                # engine-dead check still fires regularly).
+                await self._response_poller.poll(timeout=min(poll_timeout_ms, 100))
         finally:
             self._pending_rpcs.discard(rpc_id)
 

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -7,6 +7,7 @@ communicating with StageDiffusionClient via ZMQ (PUSH/PULL).
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import signal
 import time
 from concurrent.futures import ThreadPoolExecutor
@@ -30,6 +31,7 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import (
     OmniMsgpackDecoder,
     OmniMsgpackEncoder,
 )
+from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -54,6 +56,56 @@ class StageDiffusionProc:
         self._engine: DiffusionEngine | None = None
         self._executor: ThreadPoolExecutor | None = None
         self._closed = False
+        # Set by ``run_loop`` to the live dispatch task dict so
+        # :attr:`queue_length` can report in-flight requests for the
+        # OmniCoordinator heartbeat hook.
+        self._active_tasks: dict[str, asyncio.Task] | None = None
+        # Set when a request-handler detects that the engine's multiproc
+        # executor has died (e.g. a worker process crashed and the executor's
+        # monitor thread closed it). Once set, ``run_loop`` breaks out so the
+        # outer ``except``/``finally`` can send ``DIFFUSION_PROC_DEAD`` and
+        # ``ReplicaStatus.DOWN``, then the subprocess exits non-zero. Without
+        # this, the run_loop would swallow per-request errors and keep
+        # serving 500s indefinitely while heartbeats still report UP.
+        self._fatal_event: asyncio.Event | None = None
+
+    @property
+    def queue_length(self) -> int:
+        """Number of in-flight diffusion requests.
+
+        Returns 0 before :meth:`run_loop` starts and after it exits.
+        """
+        tasks = self._active_tasks
+        return 0 if tasks is None else len(tasks)
+
+    def _is_executor_dead(self) -> bool:
+        """True iff the multiproc executor has been closed or marked failed.
+
+        Detects the "workers died but the diffusion proc is still pulling
+        requests" case: ``MultiprocDiffusionExecutor`` sets ``_closed = True``
+        and ``is_failed = True`` from its worker-monitor thread the moment any
+        worker process exits; every subsequent ``execute_request`` /
+        ``collective_rpc`` then raises ``RuntimeError("DiffusionExecutor is
+        closed.")`` inside the engine. Callers in ``run_loop`` use this to
+        decide whether a per-request failure is recoverable or fatal.
+        """
+        if self._engine is None:
+            return False
+        executor = getattr(self._engine, "executor", None)
+        if executor is None:
+            return False
+        return bool(getattr(executor, "_closed", False) or getattr(executor, "is_failed", False))
+
+    def _signal_fatal_engine_failure(self, reason: str) -> None:
+        """Idempotently signal ``run_loop`` to tear down on a fatal engine error."""
+        if self._fatal_event is None or self._fatal_event.is_set():
+            return
+        logger.error(
+            "[StageDiffusionProc] fatal engine failure detected (%s); "
+            "signaling run_loop to send DIFFUSION_PROC_DEAD and exit.",
+            reason,
+        )
+        self._fatal_event.set()
 
     # ------------------------------------------------------------------
     # Initialization
@@ -309,6 +361,14 @@ class StageDiffusionProc:
         decoder = OmniMsgpackDecoder()
 
         tasks: dict[str, asyncio.Task] = {}
+        # Expose the live task dict so :attr:`queue_length` (used by the
+        # OmniCoordinator heartbeat hook) can read the in-flight count.
+        self._active_tasks = tasks
+        # Wakes the main recv loop when a request-handler detects a fatal
+        # engine failure so we tear down promptly instead of swallowing
+        # "DiffusionExecutor is closed" on every subsequent request.
+        fatal_event = asyncio.Event()
+        self._fatal_event = fatal_event
 
         async def _dispatch_request(
             request_id: str,
@@ -342,12 +402,40 @@ class StageDiffusionProc:
                         }
                     )
                 )
+                # Per-request errors are usually recoverable, but a closed
+                # executor means every future request will get the same
+                # "DiffusionExecutor is closed" error. Signal the main loop
+                # to send DIFFUSION_PROC_DEAD and exit so the head's hub
+                # demotes this replica instead of waiting on the heartbeat
+                # timeout (~30 s by default).
+                if self._is_executor_dead():
+                    self._signal_fatal_engine_failure(f"add_request {request_id}: {e!s}")
             finally:
                 tasks.pop(request_id, None)
 
         try:
             while True:
-                raw = await request_socket.recv()
+                # Await recv and fatal_event concurrently so the loop wakes
+                # up immediately when a per-request handler signals a fatal
+                # engine failure — even if no fresh ZMQ frame arrives.
+                recv_task: asyncio.Task = asyncio.ensure_future(request_socket.recv())
+                fatal_task: asyncio.Task = asyncio.ensure_future(fatal_event.wait())
+                try:
+                    done, pending = await asyncio.wait(
+                        [recv_task, fatal_task],
+                        return_when=asyncio.FIRST_COMPLETED,
+                    )
+                finally:
+                    for waiter in (recv_task, fatal_task):
+                        if not waiter.done():
+                            waiter.cancel()
+                            with contextlib.suppress(asyncio.CancelledError, Exception):
+                                await waiter
+                if fatal_event.is_set():
+                    raise RuntimeError(
+                        "StageDiffusionProc executor reported permanent failure; tearing down the diffusion subprocess."
+                    )
+                raw = recv_task.result()
                 msg = decoder.decode(raw)
                 msg_type = msg.get("type")
 
@@ -397,6 +485,11 @@ class StageDiffusionProc:
                                     }
                                 )
                             )
+                            # Same rationale as the single-request path: a
+                            # closed executor turns every subsequent batch
+                            # into a 500, so escalate now.
+                            if self._is_executor_dead():
+                                self._signal_fatal_engine_failure(f"add_batch_request {rid}: {e!s}")
                         finally:
                             tasks.pop(rid, None)
 
@@ -446,6 +539,13 @@ class StageDiffusionProc:
                                 }
                             )
                         )
+                        # Collective RPCs run through the same multiproc
+                        # executor — a closed executor means every future
+                        # RPC fails the same way, so tear down promptly.
+                        if self._is_executor_dead():
+                            self._signal_fatal_engine_failure(
+                                f"collective_rpc {msg['method']} (rpc_id={rpc_id}): {e!s}"
+                            )
 
                 elif msg_type == "shutdown":
                     break
@@ -466,6 +566,8 @@ class StageDiffusionProc:
             if tasks:
                 await asyncio.gather(*tasks.values(), return_exceptions=True)
 
+            self._active_tasks = None
+            self._fatal_event = None
             request_socket.close()
             response_socket.close()
             ctx.term()
@@ -504,8 +606,22 @@ class StageDiffusionProc:
         handshake_address: str,
         request_address: str,
         response_address: str,
+        *,
+        omni_coordinator_address: str | None = None,
+        omni_stage_id: int | None = None,
+        omni_replica_id: int = 0,
     ) -> None:
-        """Entry point for the diffusion subprocess."""
+        """Entry point for the diffusion subprocess.
+
+        Omni-specific kwargs (mirroring :meth:`StageEngineCoreProc.run_stage_core`):
+          - ``omni_coordinator_address``: ROUTER address of the head-side
+            OmniCoordinator. When set, a :class:`OmniCoordClientForStage`
+            reports the diffusion replica's status + queue length.
+          - ``omni_stage_id``: logical stage id; required when
+            ``omni_coordinator_address`` is set.
+          - ``omni_replica_id``: cluster-unique replica id within the
+            stage (logging / metrics only).
+        """
         shutdown_requested = False
 
         def signal_handler(signum: int, frame: Any) -> None:
@@ -518,6 +634,7 @@ class StageDiffusionProc:
         signal.signal(signal.SIGINT, signal_handler)
 
         proc = cls(model, od_config)
+        coord_client: OmniCoordClientForStage | None = None
         try:
             proc.initialize()
 
@@ -529,6 +646,32 @@ class StageDiffusionProc:
             handshake_socket.close()
             handshake_ctx.term()
 
+            # Wire OmniCoordClientForStage *after* READY so that the head
+            # has bound its head-side request/response sockets — the
+            # address pair we report is the same pair this proc binds to
+            # (request/response addresses passed in).
+            if omni_coordinator_address is not None:
+                if omni_stage_id is None:
+                    raise ValueError("omni_stage_id must be provided when omni_coordinator_address is set")
+                coord_client = OmniCoordClientForStage(
+                    coord_zmq_addr=omni_coordinator_address,
+                    input_addr=request_address,
+                    output_addr=response_address,
+                    stage_id=int(omni_stage_id),
+                )
+
+                def _refresh_queue_length() -> None:
+                    coord_client._queue_length = proc.queue_length  # type: ignore[union-attr]
+
+                coord_client._on_heartbeat = _refresh_queue_length
+
+                logger.info(
+                    "StageDiffusionProc registered with OmniCoordinator (stage_id=%d replica_id=%d coord=%s)",
+                    omni_stage_id,
+                    omni_replica_id,
+                    omni_coordinator_address,
+                )
+
             # Run async event loop
             asyncio.run(proc.run_loop(request_address, response_address))
 
@@ -539,6 +682,9 @@ class StageDiffusionProc:
             logger.exception("StageDiffusionProc encountered a fatal error.")
             raise
         finally:
+            if coord_client is not None:
+                with contextlib.suppress(RuntimeError):
+                    coord_client.close()
             proc.close()
 
 
@@ -551,10 +697,17 @@ def spawn_diffusion_proc(
     handshake_address: str | None = None,
     request_address: str | None = None,
     response_address: str | None = None,
+    *,
+    omni_coordinator_address: str | None = None,
+    omni_stage_id: int | None = None,
+    omni_replica_id: int = 0,
 ) -> tuple[BaseProcess, str, str, str]:
     """Spawn a StageDiffusionProc subprocess.
 
     Returns ``(proc, handshake_address, request_address, response_address)``.
+
+    Pass ``omni_coordinator_address`` / ``omni_stage_id`` / ``omni_replica_id``
+    to have the subprocess publish heartbeats to an OmniCoordinator.
     """
     handshake_address = handshake_address or get_open_zmq_ipc_path()
     request_address = request_address or get_open_zmq_ipc_path()
@@ -570,6 +723,9 @@ def spawn_diffusion_proc(
             "handshake_address": handshake_address,
             "request_address": request_address,
             "response_address": response_address,
+            "omni_coordinator_address": omni_coordinator_address,
+            "omni_stage_id": omni_stage_id,
+            "omni_replica_id": omni_replica_id,
         },
     )
     proc.start()

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -185,6 +185,10 @@ class DiffusionWorker:
         self.model_runner: DiffusionModelRunner | None = None
         self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
         self.lora_manager: DiffusionLoRAManager | None = None
+        # Worker-side cache of (lora_request, lora_scale) per scheduled
+        # request id. Used by step mode to recover LoRA identity for cached
+        # requests, which only carry their sched_req_id in subsequent ticks.
+        self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {}
         self.stage_id = getattr(od_config, "stage_id", 0)
         self.init_device()
         # Create model runner using the platform-specified class
@@ -367,13 +371,7 @@ class DiffusionWorker:
     def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
         """Execute one diffusion step by delegating to the model runner."""
         assert self.model_runner is not None, "Model runner not initialized"
-        if self.lora_manager is not None:
-            # Step mode does not support LoRA yet. Clear any previously active
-            # adapter first so worker-local LoRA state cannot leak in.
-            self.lora_manager.set_active_adapter(None)
-
-        if any(new_req.req.sampling_params.lora_request is not None for new_req in scheduler_output.scheduled_new_reqs):
-            raise ValueError("Step mode does not support LoRA yet.")
+        self._activate_step_lora(scheduler_output)
         profiler = self._get_profiler()
         ctx = profiler.annotate_context_manager("diffusion_step") if profiler else nullcontext()
         with ctx:
@@ -382,6 +380,43 @@ class DiffusionWorker:
             profiler.step()
         return output
 
+    def _activate_step_lora(self, scheduler_output: DiffusionSchedulerOutput) -> None:
+        """Activate the LoRA adapter for the scheduled step batch.
+
+        Newly scheduled requests register their (lora_request, lora_scale)
+        in ``_step_lora_state`` so cached requests can resolve to the same
+        adapter on later ticks. Finished requests are evicted. Batch
+        homogeneity is enforced by the scheduler via ``SamplingParamsKey``,
+        so any scheduled request id resolves to the active LoRA identity.
+        """
+        for sched_req_id in scheduler_output.finished_req_ids:
+            self._step_lora_state.pop(sched_req_id, None)
+
+        for new_req in scheduler_output.scheduled_new_reqs:
+            sampling = new_req.req.sampling_params
+            self._step_lora_state[new_req.sched_req_id] = (
+                sampling.lora_request,
+                sampling.lora_scale,
+            )
+
+        if self.lora_manager is None:
+            return
+
+        lora_request: LoRARequest | None = None
+        lora_scale = 1.0
+        for sched_req_id in scheduler_output.scheduled_req_ids:
+            entry = self._step_lora_state.get(sched_req_id)
+            if entry is not None:
+                lora_request, lora_scale = entry
+                break
+
+        try:
+            self.lora_manager.set_active_adapter(lora_request, lora_scale)
+        except Exception as exc:
+            if lora_request is not None:
+                raise
+            logger.warning("LoRA activation skipped: %s", exc)
+
     def load_weights(self, weights) -> set[str]:
         """Load weights by delegating to the model runner."""
         assert self.model_runner is not None, "Model runner not initialized"