feat: support pipefusion for flux2

lightx2v/common/distributed/__init__.py

@@ -0,0 +1,11 @@
+from .pipeline_comm import PipelineComm
+from .pipeline_state import (
+    PipelineRuntimeState,
+    get_pipeline_parallel_rank,
+    get_pipeline_parallel_world_size,
+    get_pipeline_runtime_state,
+    get_pp_group,
+    init_pipeline_parallel_state,
+    is_pipeline_first_stage,
+    is_pipeline_last_stage,
+)

lightx2v/common/distributed/pipeline_comm.py

@@ -0,0 +1,115 @@
+"""P2P communication manager for pipeline parallelism.
+
+Uses the default CUDA stream for all NCCL operations. NCCL's internal
+stream handles the actual async data transfer, so a separate comm_stream
+is unnecessary and would only introduce cross-stream sync overhead.
+
+isend requests are returned to the caller who MUST store them to
+prevent tensor GC before the send completes.
+"""
+
+from typing import Dict, List, Tuple
+
+import torch
+import torch.distributed as dist
+
+
+class PipelineComm:
+    """P2P communication between adjacent pipeline stages."""
+
+    def __init__(self, pp_group: dist.ProcessGroup):
+        self.pp_group = pp_group
+        self.rank = dist.get_rank(pp_group)
+        self.world_size = dist.get_world_size(pp_group)
+
+        self.ranks = list(dist.get_process_group_ranks(pp_group))
+        self.prev_rank = self.ranks[(self.rank - 1) % self.world_size]
+        self.next_rank = self.ranks[(self.rank + 1) % self.world_size]
+        self._device_group = pp_group
+
+        self._recv_tasks_queue: List[Tuple[str, int]] = []
+        self._receiving_tasks: List[Tuple[object, str, int]] = []
+        self._recv_buffers: Dict[Tuple[str, int], torch.Tensor] = {}
+
+    # ------------------------------------------------------------------
+    # Synchronous send / recv (sync pipeline)
+    # ------------------------------------------------------------------
+
+    def pipeline_send(self, tensor: torch.Tensor, name: str = "latent", skip_shape: bool = False):
+        tensor = tensor.contiguous()
+        if not skip_shape:
+            shape_info = torch.tensor(
+                [tensor.ndim] + list(tensor.shape),
+                device=tensor.device,
+                dtype=torch.int64,
+            )
+            padded = torch.zeros(9, device=tensor.device, dtype=torch.int64)
+            padded[: len(shape_info)] = shape_info
+            dist.send(padded, dst=self.next_rank, group=self._device_group)
+        dist.send(tensor, dst=self.next_rank, group=self._device_group)
+
+    def pipeline_recv(self, name: str = "latent", shape=None, dtype=None) -> torch.Tensor:
+        if dtype is None:
+            dtype = torch.bfloat16
+        if shape is not None:
+            buf = torch.empty(shape, dtype=dtype, device=torch.cuda.current_device())
+            dist.recv(buf, src=self.prev_rank, group=self._device_group)
+            return buf
+        shape_info = torch.zeros(9, device=torch.cuda.current_device(), dtype=torch.int64)
+        dist.recv(shape_info, src=self.prev_rank, group=self._device_group)
+        ndim = shape_info[0].item()
+        recv_shape = tuple(shape_info[1 : 1 + ndim].tolist())
+        buf = torch.empty(recv_shape, dtype=dtype, device=torch.cuda.current_device())
+        dist.recv(buf, src=self.prev_rank, group=self._device_group)
+        return buf
+
+    # ------------------------------------------------------------------
+    # Asynchronous send / recv (async pipeline)
+    # All on default stream — NCCL's internal stream handles the real
+    # async transfer, and req.wait() just inserts a stream-side
+    # dependency (non-blocking on CPU in the same-stream case).
+    # ------------------------------------------------------------------
+
+    def pipeline_isend(self, tensor: torch.Tensor, name: str = "latent", segment_idx: int = 0):
+        """Non-blocking send on the current (default) stream.
+
+        Returns a Work object that the caller SHOULD store to prevent
+        the tensor from being garbage-collected before the send
+        completes. The Work's wait() only inserts a stream-side
+        dependency — it does not block the CPU thread.
+        """
+        tensor = tensor.contiguous()
+        return dist.isend(tensor, dst=self.next_rank, group=self._device_group)
+
+    def add_pipeline_recv_task(self, idx: int = 0, name: str = "latent", shape=None, dtype=None):
+        self._recv_tasks_queue.append((name, idx))
+        if (name, idx) not in self._recv_buffers:
+            assert shape is not None and dtype is not None
+            self._recv_buffers[(name, idx)] = torch.empty(shape, dtype=dtype, device=torch.cuda.current_device())
+
+    def recv_next(self):
+        """Post next irecv on the current (default) stream.
+
+        Non-blocking on CPU: dist.irecv enqueues the work on NCCL's
+        internal stream and returns immediately.
+        """
+        if not self._recv_tasks_queue:
+            raise ValueError("No more tasks to receive")
+        name, idx = self._recv_tasks_queue.pop(0)
+        buf = self._recv_buffers.get((name, idx))
+        assert buf is not None
+        req = dist.irecv(buf, src=self.prev_rank, group=self._device_group)
+        self._receiving_tasks.append((req, name, idx))
+
+    def get_pipeline_recv_data(self, idx: int = 0, name: str = "latent") -> torch.Tensor:
+        """Wait for and return a previously posted async receive.
+
+        In the single-stream model, req.wait() inserts a stream-side
+        wait for the NCCL op's completion on the current stream. It
+        does NOT block the CPU thread unless a timeout is set.
+        """
+        assert self._receiving_tasks
+        req, rname, ridx = self._receiving_tasks.pop(0)
+        assert rname == name and ridx == idx
+        req.wait()
+        return self._recv_buffers[(name, idx)]

lightx2v/common/distributed/pipeline_state.py

@@ -0,0 +1,152 @@
+"""Pipeline parallel runtime state and stage helpers for LightX2V.
+
+Manages patch metadata (how the latent image is split across pipeline patches)
+and provides stage-identification utilities (is_pipeline_first_stage, etc.).
+"""
+
+from typing import List, Optional
+
+import torch.distributed as dist
+
+# ---------------------------------------------------------------------------
+# Global state
+# ---------------------------------------------------------------------------
+
+_pp_group: Optional[dist.ProcessGroup] = None
+_runtime_state: Optional["PipelineRuntimeState"] = None
+
+
+def init_pipeline_parallel_state(pp_group: dist.ProcessGroup):
+    """Register the pipeline-parallel process group.  Called once during
+    ``set_parallel_config`` when ``pp_size > 1``."""
+    global _pp_group, _runtime_state
+    _pp_group = pp_group
+    _runtime_state = PipelineRuntimeState()
+
+
+# ---------------------------------------------------------------------------
+# Stage helpers
+# ---------------------------------------------------------------------------
+
+
+def get_pp_group() -> dist.ProcessGroup:
+    assert _pp_group is not None, "pipeline parallel group is not initialised"
+    return _pp_group
+
+
+def get_pipeline_parallel_rank() -> int:
+    if _pp_group is None:
+        return 0
+    return dist.get_rank(_pp_group)
+
+
+def get_pipeline_parallel_world_size() -> int:
+    if _pp_group is None:
+        return 1
+    return dist.get_world_size(_pp_group)
+
+
+def is_pipeline_first_stage() -> bool:
+    return get_pipeline_parallel_rank() == 0
+
+
+def is_pipeline_last_stage() -> bool:
+    return get_pipeline_parallel_rank() == get_pipeline_parallel_world_size() - 1
+
+
+def get_pipeline_runtime_state() -> "PipelineRuntimeState":
+    assert _runtime_state is not None, "PipelineRuntimeState not initialised"
+    return _runtime_state
+
+
+# ---------------------------------------------------------------------------
+# Runtime state
+# ---------------------------------------------------------------------------
+
+
+class PipelineRuntimeState:
+    """Runtime metadata for patch-level pipeline parallelism (PipeFusion).
+
+    Computes how the latent token sequence is split into *patches* so that
+    each pipeline stage processes a subset of patches in async mode.
+    """
+
+    def __init__(self):
+        self.num_pipeline_patch: int = 1
+        self.pipeline_patch_idx: int = 0
+        self.patch_mode: bool = False  # True = async, False = sync
+        self.warmup_steps: int = 1
+
+        # Patch metadata (along the latent token / sequence dimension)
+        self.pp_patches_token_num: List[int] = [0]
+        self.pp_patches_token_start_end_idx_global: List[List[int]] = [[0, 0]]
+
+        # Input parameters
+        self.height: int = 0
+        self.width: int = 0
+        self.batch_size: int = 1
+        self.packed_h: int = 0
+        self.packed_w: int = 0
+        self.vae_scale_factor: int = 16
+        self.patch_size: int = 1
+
+    # -- configuration -------------------------------------------------------
+
+    def set_input_parameters(
+        self,
+        height: int,
+        width: int,
+        batch_size: int = 1,
+        num_pipeline_patch: Optional[int] = None,
+        warmup_steps: int = 1,
+        vae_scale_factor: int = 16,
+        patch_size: int = 1,
+        total_tokens: Optional[int] = None,
+    ):
+        self.height = height
+        self.width = width
+        self.batch_size = batch_size
+        self.vae_scale_factor = vae_scale_factor
+        self.patch_size = patch_size
+        self.warmup_steps = warmup_steps
+        if num_pipeline_patch is not None:
+            self.num_pipeline_patch = num_pipeline_patch
+
+        if total_tokens is not None:
+            self.packed_h = 0
+            self.packed_w = 0
+            tok_count = total_tokens
+        else:
+            # Compute packed dimensions
+            multiple_of = vae_scale_factor * 2
+            self.packed_h = height // multiple_of
+            self.packed_w = width // multiple_of
+            tok_count = self.packed_h * self.packed_w
+
+        # Split tokens evenly across patches
+        base = tok_count // self.num_pipeline_patch
+        remainder = tok_count % self.num_pipeline_patch
+        self.pp_patches_token_num = []
+        self.pp_patches_token_start_end_idx_global = []
+        start = 0
+        for i in range(self.num_pipeline_patch):
+            n = base + (1 if i < remainder else 0)
+            self.pp_patches_token_num.append(n)
+            self.pp_patches_token_start_end_idx_global.append([start, start + n])
+            start += n
+
+    # -- patch mode ----------------------------------------------------------
+
+    def set_patched_mode(self, patch_mode: bool):
+        self.patch_mode = patch_mode
+        self.pipeline_patch_idx = 0
+
+    def next_patch(self):
+        if self.patch_mode:
+            self.pipeline_patch_idx += 1
+            if self.pipeline_patch_idx >= self.num_pipeline_patch:
+                self.pipeline_patch_idx = 0
+
+    @property
+    def current_patch_token_start_end(self) -> List[int]:
+        return self.pp_patches_token_start_end_idx_global[self.pipeline_patch_idx]

lightx2v/common/ops/attn/flash_attn.py

@@ -52,9 +52,10 @@ def apply(
         softmax_scale = kwargs.get("softmax_scale", None)
         if len(q.shape) == 3:
             bs = 1
+            total_seqlen = q.shape[0]
         elif len(q.shape) == 4:
             bs = q.shape[0]
-        total_seqlen = bs * max_seqlen_q
+            total_seqlen = bs * q.shape[1]
 
         if bs == 1:
             if len(q.shape) == 3:
@@ -127,9 +128,10 @@ def apply(
         softmax_scale = kwargs.get("softmax_scale", None)
         if len(q.shape) == 3:
             bs = 1
+            total_seqlen = q.shape[0]
         elif len(q.shape) == 4:
             bs = q.shape[0]
-        total_seqlen = bs * max_seqlen_q
+            total_seqlen = bs * q.shape[1]
 
         if bs == 1:
             if len(q.shape) == 3:
@@ -209,7 +211,7 @@ def apply(
             k,
             v,
         )
-        x = x.reshape(bs * max_seqlen_q, -1)
+        x = x.reshape(q.shape[1], -1)
         return x
 
 

lightx2v/models/networks/base_model.py

@@ -399,7 +399,11 @@ def _load_safetensor_to_dict(self, file_path, unified_dtype, sensitive_layer):
         remove_keys = self.remove_keys if hasattr(self, "remove_keys") else []
         preserve_keys = self.preserved_keys if hasattr(self, "preserved_keys") else None  # None means all keys are preserved, otherwise only keys in preserve_keys are preserved
 
-        if self.device.type != "cpu" and dist.is_initialized():
+        # In PipeFusion mode, load weights to CPU first to avoid OOM — each
+        # stage only needs a subset of block weights on GPU.
+        if self.config.get("pipefusion_parallel", False):
+            device = "cpu"
+        elif self.device.type != "cpu" and dist.is_initialized():
             device = dist.get_rank()
         else:
             device = str(self.device)

lightx2v/models/networks/flux2/infer/pipefusion/__init__.py

@@ -0,0 +1,2 @@
+from .pipeline_driver import Flux2PipelineDriver
+from .transformer_infer import Flux2PipeFusionTransformerInfer