vllm-omni/vllm_omni/diffusion/distributed/pipeline_parallel.py

# 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