From 50e8996318f31b2affbfe555ce83c6fe8fe0745a Mon Sep 17 00:00:00 2001
From: Qin-sx <Qin_sx@outlook.com>
Date: Sun, 19 Jul 2026 23:54:11 +0800
Subject: [PATCH 01/12] feat: support pipefusion for flux2

	new file:   lightx2v/common/distributed/__init__.py
	new file:   lightx2v/common/distributed/pipeline_comm.py
	new file:   lightx2v/common/distributed/pipeline_state.py
	modified:   lightx2v/models/networks/base_model.py
	new file:   lightx2v/models/networks/flux2/infer/pipefusion/__init__.py
	new file:   lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
	new file:   lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
	modified:   lightx2v/models/networks/flux2/infer/pre_infer.py
	modified:   lightx2v/models/networks/flux2/infer/transformer_infer.py
	modified:   lightx2v/models/networks/flux2/model.py
	modified:   lightx2v/models/networks/flux2/weights/transformer_weights.py
	modified:   lightx2v/models/runners/flux2/flux2_runner.py
	modified:   lightx2v/models/schedulers/flux2/scheduler.py
	modified:   lightx2v/pipeline.py
	modified:   lightx2v/utils/set_config.py
---
 lightx2v/common/distributed/__init__.py       |  11 +
 lightx2v/common/distributed/pipeline_comm.py  | 115 +++++
 lightx2v/common/distributed/pipeline_state.py | 146 +++++++
 lightx2v/models/networks/base_model.py        |   6 +-
 .../flux2/infer/pipefusion/__init__.py        |   2 +
 .../flux2/infer/pipefusion/pipeline_driver.py | 395 ++++++++++++++++++
 .../infer/pipefusion/transformer_infer.py     | 200 +++++++++
 .../models/networks/flux2/infer/pre_infer.py  |  43 ++
 .../networks/flux2/infer/transformer_infer.py |  39 +-
 lightx2v/models/networks/flux2/model.py       |  12 +-
 .../flux2/weights/transformer_weights.py      |  46 +-
 lightx2v/models/runners/flux2/flux2_runner.py | 102 ++++-
 lightx2v/models/schedulers/flux2/scheduler.py |   9 +
 lightx2v/pipeline.py                          |   5 +-
 lightx2v/utils/set_config.py                  |  20 +-
 15 files changed, 1128 insertions(+), 23 deletions(-)
 create mode 100644 lightx2v/common/distributed/__init__.py
 create mode 100644 lightx2v/common/distributed/pipeline_comm.py
 create mode 100644 lightx2v/common/distributed/pipeline_state.py
 create mode 100644 lightx2v/models/networks/flux2/infer/pipefusion/__init__.py
 create mode 100644 lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
 create mode 100644 lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py

diff --git a/lightx2v/common/distributed/__init__.py b/lightx2v/common/distributed/__init__.py
new file mode 100644
index 000000000..0f979d8f9
--- /dev/null
+++ b/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,
+)
diff --git a/lightx2v/common/distributed/pipeline_comm.py b/lightx2v/common/distributed/pipeline_comm.py
new file mode 100644
index 000000000..521913272
--- /dev/null
+++ b/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)]
diff --git a/lightx2v/common/distributed/pipeline_state.py b/lightx2v/common/distributed/pipeline_state.py
new file mode 100644
index 000000000..9838f7ce7
--- /dev/null
+++ b/lightx2v/common/distributed/pipeline_state.py
@@ -0,0 +1,146 @@
+"""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,
+    ):
+        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
+
+        # Compute packed dimensions (matching Flux2Runner.set_target_shape)
+        multiple_of = vae_scale_factor * 2
+        self.packed_h = height // multiple_of
+        self.packed_w = width // multiple_of
+        total_tokens = self.packed_h * self.packed_w
+
+        # Split tokens evenly across patches
+        base = total_tokens // self.num_pipeline_patch
+        remainder = total_tokens % 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]
diff --git a/lightx2v/models/networks/base_model.py b/lightx2v/models/networks/base_model.py
index 4145a3ec1..eb2ee5d14 100644
--- a/lightx2v/models/networks/base_model.py
+++ b/lightx2v/models/networks/base_model.py
@@ -389,7 +389,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)
diff --git a/lightx2v/models/networks/flux2/infer/pipefusion/__init__.py b/lightx2v/models/networks/flux2/infer/pipefusion/__init__.py
new file mode 100644
index 000000000..fbd3b4694
--- /dev/null
+++ b/lightx2v/models/networks/flux2/infer/pipefusion/__init__.py
@@ -0,0 +1,2 @@
+from .pipeline_driver import Flux2PipelineDriver
+from .transformer_infer import Flux2PipeFusionTransformerInfer
diff --git a/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py b/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
new file mode 100644
index 000000000..d60d47220
--- /dev/null
+++ b/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
@@ -0,0 +1,395 @@
+"""Sync/Async pipeline driver for Flux2 PipeFusion.
+
+Orchestrates the denoising loop across pipeline stages:
+- **Sync pipeline** (warmup): each timestep, all stages process the full latent
+  sequentially (stage 0 -> stage 1 -> ... -> last stage).
+- **Async pipeline** (main loop): each timestep, the latent is split into
+  patches; stages process different patches concurrently, overlapping compute
+  and P2P communication.
+"""
+
+import torch
+import torch.distributed as dist
+
+from lightx2v.common.distributed import (
+    PipelineComm,
+    get_pipeline_parallel_world_size,
+    get_pipeline_runtime_state,
+    get_pp_group,
+    is_pipeline_first_stage,
+    is_pipeline_last_stage,
+)
+
+
+class Flux2PipelineDriver:
+    """Drives the PipeFusion denoising loop for Flux2."""
+
+    def __init__(self, model, config):
+        self.model = model
+        self.config = config
+        self.state = get_pipeline_runtime_state()
+        self.pp_comm = PipelineComm(get_pp_group())
+        self._is_first = is_pipeline_first_stage()
+        self._is_last = is_pipeline_last_stage()
+        self._pp_world_size = get_pipeline_parallel_world_size()
+        self._dtype = config.get("dtype", torch.bfloat16)
+        if isinstance(self._dtype, str):
+            self._dtype = getattr(torch, self._dtype)
+
+    # ==================================================================
+    # Public entry point
+    # ==================================================================
+
+    def run_pipeline(self, latents, prompt_embeds, text_ids, latent_image_ids, timesteps, scheduler, do_cfg=False, negative_prompt_embeds=None, negative_text_ids=None):
+        """Run the full denoising loop with PipeFusion.
+
+        Returns final latents on the last stage, ``None`` on other stages.
+        """
+        warmup_steps = self.state.warmup_steps
+
+        if self._pp_world_size > 1 and len(timesteps) > warmup_steps:
+            latents = self._sync_pipeline(
+                latents,
+                prompt_embeds,
+                text_ids,
+                latent_image_ids,
+                timesteps[:warmup_steps],
+                scheduler,
+                do_cfg=do_cfg,
+                negative_prompt_embeds=negative_prompt_embeds,
+                negative_text_ids=negative_text_ids,
+            )
+            latents = self._async_pipeline(
+                latents,
+                prompt_embeds,
+                text_ids,
+                latent_image_ids,
+                timesteps[warmup_steps:],
+                scheduler,
+                do_cfg=do_cfg,
+                negative_prompt_embeds=negative_prompt_embeds,
+                negative_text_ids=negative_text_ids,
+            )
+        else:
+            latents = self._sync_pipeline(
+                latents,
+                prompt_embeds,
+                text_ids,
+                latent_image_ids,
+                timesteps,
+                scheduler,
+                do_cfg=do_cfg,
+                negative_prompt_embeds=negative_prompt_embeds,
+                negative_text_ids=negative_text_ids,
+            )
+        return latents
+
+    # ==================================================================
+    # Sync pipeline (warmup)
+    # ==================================================================
+
+    def _sync_pipeline(self, latents, prompt_embeds, text_ids, latent_image_ids, timesteps, scheduler, do_cfg=False, negative_prompt_embeds=None, negative_text_ids=None):
+        self.state.set_patched_mode(patch_mode=False)
+
+        for step_idx, t in enumerate(timesteps):
+            scheduler.step_index = step_idx
+            scheduler.step_pre(step_idx)
+
+            if do_cfg:
+                # Conditional pass
+                cond_result = self._sync_pass(
+                    latents,
+                    prompt_embeds,
+                    text_ids,
+                    latent_image_ids,
+                    t,
+                    scheduler,
+                )
+                # Unconditional pass
+                uncond_result = self._sync_pass(
+                    latents,
+                    negative_prompt_embeds,
+                    negative_text_ids or text_ids,
+                    latent_image_ids,
+                    t,
+                    scheduler,
+                )
+                if self._is_last:
+                    noise_pred_cond = cond_result
+                    noise_pred_uncond = uncond_result
+                    guidance_scale = self.config.get("sample_guide_scale", 1.0)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+                    scheduler.noise_pred = noise_pred
+                    scheduler.latents = latents
+                    scheduler.step_post()
+                    latents = scheduler.latents
+            else:
+                noise_pred = self._sync_pass(
+                    latents,
+                    prompt_embeds,
+                    text_ids,
+                    latent_image_ids,
+                    t,
+                    scheduler,
+                )
+                if self._is_last:
+                    scheduler.noise_pred = noise_pred
+                    scheduler.latents = latents
+                    scheduler.step_post()
+                    latents = scheduler.latents
+
+            # P2P: last stage sends updated latents to first stage (circular)
+            # Only rank 0 needs updated latents (for x_embedder in next step).
+            # Ranks 1-6 don't participate — no global sync barrier.
+            if self._pp_world_size > 1:
+                if self._is_last:
+                    # Last stage sends to first stage
+                    dist.send(latents.contiguous(), dst=self.pp_comm.ranks[0], group=self.pp_comm.pp_group)
+                elif self._is_first:
+                    # First stage receives from last stage
+                    latents = torch.empty_like(latents)
+                    dist.recv(latents, src=self.pp_comm.ranks[-1], group=self.pp_comm.pp_group)
+                # Ranks 1-6: no op (don't need updated latents)
+
+        return latents
+
+    def _sync_pass(self, latents, prompt_embeds, text_ids, latent_image_ids, t, scheduler):
+        """Single sync forward pass through all stages.
+
+        Returns ``noise_pred`` on last stage, ``None`` on other stages.
+        P2P always carries separate (image, text) streams.
+        """
+        # NOTE: do NOT clear KV cache here. Sync mode populates per-patch
+        # slots so async mode can use them as "stale" KV for global attention.
+
+        if self._is_first:
+            # First stage: run pre_infer + blocks
+            pre_infer_out = self.model.pre_infer.infer(
+                weights=self.model.pre_weight,
+                hidden_states=latents,
+                encoder_hidden_states=prompt_embeds,
+                txt_ids=text_ids,
+                img_ids=latent_image_ids,
+            )
+            hidden_states, enc_hidden, num_txt = self.model.transformer_infer.infer(self.model.transformer_weights, pre_infer_out)
+
+            if self._is_last:
+                return self._run_post_infer(hidden_states, enc_hidden, num_txt, pre_infer_out.timestep)
+            else:
+                # Always send both latent and encoder_hidden_state (skip_shape
+                # to avoid .item() CPU-GPU sync on receiver)
+                self.pp_comm.pipeline_send(hidden_states, name="latent", skip_shape=True)
+                self.pp_comm.pipeline_send(enc_hidden, name="encoder_hidden_state", skip_shape=True)
+                return None
+        else:
+            # Non-first stage: always receive both streams
+            # Pass pre-computed shapes to avoid .item() CPU-GPU sync
+            inner_dim = self.config.get("num_attention_heads", 24) * self.config.get("attention_head_dim", 64)
+            if latent_image_ids.ndim == 3:
+                img_len = latent_image_ids.shape[1]
+            else:
+                img_len = latent_image_ids.shape[0]
+            if prompt_embeds is not None:
+                txt_len = prompt_embeds.shape[1] if prompt_embeds.ndim == 3 else prompt_embeds.shape[0]
+            else:
+                txt_len = 0
+            hidden_states = self.pp_comm.pipeline_recv(name="latent", shape=(img_len, inner_dim), dtype=self._dtype)
+            enc_hidden = self.pp_comm.pipeline_recv(name="encoder_hidden_state", shape=(txt_len, inner_dim), dtype=self._dtype)
+
+            pre_infer_out = self.model.pre_infer.infer_partial(
+                weights=self.model.pre_weight,
+                hidden_states=hidden_states,
+                encoder_hidden_states=enc_hidden,
+                txt_ids=text_ids,
+                img_ids=latent_image_ids,
+            )
+            hidden_states, enc_hidden, num_txt = self.model.transformer_infer.infer(self.model.transformer_weights, pre_infer_out)
+
+            if self._is_last:
+                return self._run_post_infer(hidden_states, enc_hidden, num_txt, pre_infer_out.timestep)
+            else:
+                self.pp_comm.pipeline_send(hidden_states, name="latent", skip_shape=True)
+                self.pp_comm.pipeline_send(enc_hidden, name="encoder_hidden_state", skip_shape=True)
+                return None
+
+    # ==================================================================
+    # Async pipeline (main loop)
+    # ==================================================================
+
+    def _async_pipeline(self, latents, prompt_embeds, text_ids, latent_image_ids, timesteps, scheduler, do_cfg=False, negative_prompt_embeds=None, negative_text_ids=None):
+        self.state.set_patched_mode(patch_mode=True)
+        num_patch = self.state.num_pipeline_patch
+        patch_token_nums = self.state.pp_patches_token_num
+        inner_dim = self.config.get("num_attention_heads", 24) * self.config.get("attention_head_dim", 64)
+        # Raw latent channels (before x_embedder): rank=0 recv from last stage
+        # gets raw latents [1, L, C_in]; other stages recv embedded [L, D]
+        raw_channels = self.config.get("transformer_in_channels", self.config.get("in_channels", 128))
+
+        # Split latents into patches (dim=1 for [B, L, C])
+        if self._is_first or self._is_last:
+            patch_latents = list(latents.split(patch_token_nums, dim=1))
+        else:
+            patch_latents = [None] * num_patch
+
+        # Split image ids by patch
+        patch_latent_image_ids = []
+        for start, end in self.state.pp_patches_token_start_end_idx_global:
+            if latent_image_ids.ndim == 3:
+                patch_latent_image_ids.append(latent_image_ids[:, start:end, :])
+            else:
+                patch_latent_image_ids.append(latent_image_ids[start:end, :])
+
+        # Compute txt_len for buffer allocation
+        if prompt_embeds is not None:
+            txt_len = prompt_embeds.shape[1] if prompt_embeds.ndim == 3 else prompt_embeds.shape[0]
+        else:
+            txt_len = 0
+
+        # Pre-allocate recv buffers and pre-post all receives
+        # First stage: receives raw latents [1, L, C_in] from last stage (circular)
+        # Non-first stages: receives embedded encoder + latent [L, D] from previous stage
+        recv_timesteps = len(timesteps) - 1 if self._is_first else len(timesteps)
+        for _ in range(recv_timesteps):
+            if not self._is_first:
+                self.pp_comm.add_pipeline_recv_task(
+                    0,
+                    "encoder_hidden_state",
+                    shape=(txt_len, inner_dim),
+                    dtype=self._dtype,
+                )
+            for patch_idx in range(num_patch):
+                # First stage (rank=0) receives raw latents [1, L, C_in] from
+                # last stage; other stages receive embedded [L, D]
+                if self._is_first:
+                    latent_shape = (1, patch_token_nums[patch_idx], raw_channels)
+                else:
+                    latent_shape = (patch_token_nums[patch_idx], inner_dim)
+                self.pp_comm.add_pipeline_recv_task(
+                    patch_idx,
+                    "latent",
+                    shape=latent_shape,
+                    dtype=self._dtype,
+                )
+
+        last_patch_latents = [None] * num_patch if self._is_last else None
+        first_async_recv = True
+        total_steps = len(timesteps)
+
+        # Track pending isend requests to prevent tensor GC before send completes
+        pending_isends = []
+
+        for i, t in enumerate(timesteps):
+            scheduler.step_index = i + self.state.warmup_steps
+            scheduler.step_pre(scheduler.step_index)
+
+            for patch_idx in range(num_patch):
+                if self._is_last:
+                    last_patch_latents[patch_idx] = patch_latents[patch_idx]
+
+                # ---- 1. Receive current patch's data ----
+                if self._is_first and i == 0:
+                    pass  # first stage, first step: has initial latents
+                else:
+                    if first_async_recv:
+                        if not self._is_first and patch_idx == 0:
+                            self.pp_comm.recv_next()
+                        self.pp_comm.recv_next()
+                        first_async_recv = False
+                    if not self._is_first and patch_idx == 0:
+                        last_encoder_hidden_states = self.pp_comm.get_pipeline_recv_data(0, "encoder_hidden_state")
+                    if not (self._is_first and i == 0):
+                        patch_latents[patch_idx] = self.pp_comm.get_pipeline_recv_data(patch_idx, "latent")
+
+                # ---- 2. Compute (default stream) ----
+                cur_enc = prompt_embeds if self._is_first else last_encoder_hidden_states
+                result = self._async_backbone(
+                    patch_latents[patch_idx],
+                    cur_enc,
+                    text_ids,
+                    patch_latent_image_ids[patch_idx],
+                    scheduler,
+                )
+
+                # ---- 3. Send result (default stream, after compute) ----
+                # Store isend request to prevent tensor GC before send completes
+                if self._is_last:
+                    noise_pred = result
+                    scheduler.scheduler._step_index = i + self.state.warmup_steps
+                    patch_latents[patch_idx] = scheduler.scheduler.step(noise_pred, t, last_patch_latents[patch_idx], return_dict=False)[0]
+                    if i != total_steps - 1:
+                        req = self.pp_comm.pipeline_isend(patch_latents[patch_idx], name="latent", segment_idx=patch_idx)
+                        pending_isends.append((req, patch_latents[patch_idx]))
+                else:
+                    hidden_states, next_enc = result
+                    if patch_idx == 0:
+                        req = self.pp_comm.pipeline_isend(next_enc, name="encoder_hidden_state", segment_idx=0)
+                        pending_isends.append((req, next_enc))
+                    req = self.pp_comm.pipeline_isend(hidden_states, name="latent", segment_idx=patch_idx)
+                    pending_isends.append((req, hidden_states))
+
+                # ---- 4. Post next irecv (default stream — NCCL internal
+                # stream handles the actual async transfer; no cross-stream
+                # sync needed.) ----
+                if not (self._is_first and i == 0):
+                    is_last_step = i == total_steps - 1
+                    is_last_patch = patch_idx == num_patch - 1
+                    if not (is_last_step and is_last_patch):
+                        if self._is_first:
+                            self.pp_comm.recv_next()
+                        else:
+                            if is_last_patch:
+                                self.pp_comm.recv_next()
+                            self.pp_comm.recv_next()
+
+                # ---- 5. Wait for old isends (limit pending to prevent GC issues) ----
+                while len(pending_isends) > 4:
+                    old_req, _ = pending_isends.pop(0)
+                    old_req.wait()
+
+                self.state.next_patch()
+
+        # Wait for all remaining isends before returning
+        for req, _ in pending_isends:
+            req.wait()
+        pending_isends.clear()
+
+        if self._is_last:
+            return torch.cat(patch_latents, dim=1)
+        return None
+
+    def _async_backbone(self, patch_latent, encoder_hidden_states, text_ids, patch_img_ids, scheduler):
+        """Backbone forward for a single patch in async mode."""
+        if self._is_first:
+            pre_infer_out = self.model.pre_infer.infer(
+                weights=self.model.pre_weight,
+                hidden_states=patch_latent,
+                encoder_hidden_states=encoder_hidden_states,
+                txt_ids=text_ids,
+                img_ids=patch_img_ids,
+            )
+        else:
+            pre_infer_out = self.model.pre_infer.infer_partial(
+                weights=self.model.pre_weight,
+                hidden_states=patch_latent,
+                encoder_hidden_states=encoder_hidden_states,
+                txt_ids=text_ids,
+                img_ids=patch_img_ids,
+            )
+
+        hidden_states, enc_hidden, num_txt = self.model.transformer_infer.infer(self.model.transformer_weights, pre_infer_out)
+
+        if self._is_last:
+            return self._run_post_infer(hidden_states, enc_hidden, num_txt, pre_infer_out.timestep)
+        else:
+            return (hidden_states, enc_hidden)
+
+    # ==================================================================
+    # Shared helpers
+    # ==================================================================
+
+    def _run_post_infer(self, hidden_states, enc_hidden, num_txt, timestep):
+        """Run post_infer on the last stage."""
+        if enc_hidden is None and num_txt > 0:
+            hidden_states = hidden_states[num_txt:, ...]
+        noise_pred = self.model.post_infer.infer(self.model.post_weight, hidden_states, timestep)
+        return noise_pred
diff --git a/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py b/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
new file mode 100644
index 000000000..5c1bbb684
--- /dev/null
+++ b/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
@@ -0,0 +1,200 @@
+"""PipeFusion-enabled transformer infer for Flux2.
+
+Subclasses ``Flux2TransformerInfer`` to:
+1. Run only the current pipeline stage's block subset.
+2. Apply stale-KV caching in async (patched) mode: image KV is cached across
+   patches while text KV stays fresh.
+3. Return ``(hidden_states, encoder_hidden_states, num_txt_tokens)`` so the
+   pipeline driver can P2P-pass intermediate activations between stages.
+"""
+
+import torch
+import torch.nn.functional as F
+
+from ..transformer_infer import Flux2TransformerInfer
+
+
+class Flux2PipeFusionTransformerInfer(Flux2TransformerInfer):
+    """Transformer infer with PipeFusion block splitting and stale-KV cache."""
+
+    def __init__(self, config):
+        super().__init__(config)
+        from lightx2v.common.distributed import (
+            get_pipeline_runtime_state,
+            is_pipeline_first_stage,
+            is_pipeline_last_stage,
+        )
+
+        self.pipeline_state = get_pipeline_runtime_state()
+        self._is_first_stage = is_pipeline_first_stage()
+        self._is_last_stage = is_pipeline_last_stage()
+
+        # Stale-KV cache: block_idx -> [ [k, v] per patch slot ]
+        self._kv_cache: dict = {}
+
+        # Pre-allocated full K/V buffers (lazily created on first async use).
+        # Avoids repeated torch.cat allocations per timestep.
+        self._full_k_buf = None
+        self._full_v_buf = None
+
+    # ------------------------------------------------------------------
+    # Stale-KV hook (overrides base class no-op)
+    # ------------------------------------------------------------------
+
+    def _maybe_apply_stale_kv(self, key, value, num_txt_tokens, block_idx):
+        """Per-patch-slot KV cache for PipeFusion.
+
+        Semantics:
+        - Cache is indexed by (block_idx, patch_slot). Each slot stores the
+          image K/V computed for that patch when it was last processed.
+        - SYNC mode: split full image K/V by patch, populate ALL slots.
+          Return input unchanged (full attention runs normally).
+        - ASYNC mode: update current patch's slot with fresh K/V; use full
+          cache (fresh + stale from prior timestep) for attention.
+
+        Optimization: pre-allocated buffers + copy_ instead of torch.cat
+        to avoid memory allocations per generation.
+        """
+        num_patch = self.pipeline_state.num_pipeline_patch
+        if num_patch <= 1 or num_txt_tokens <= 0:
+            return key, value
+
+        # Split text / image along sequence dim
+        text_key, img_key = key.split([num_txt_tokens, key.shape[0] - num_txt_tokens], dim=0)
+        text_value, img_value = value.split([num_txt_tokens, value.shape[0] - num_txt_tokens], dim=0)
+
+        patch_token_nums = self.pipeline_state.pp_patches_token_num
+
+        if not self.pipeline_state.patch_mode:
+            # Sync mode: split full image K/V by patch, populate all slots.
+            # .clone() ensures cached tensors own their storage (views of
+            # transient QKV would become invalid after this timestep).
+            if block_idx not in self._kv_cache:
+                self._kv_cache[block_idx] = [None] * num_patch
+            split_ks = img_key.split(patch_token_nums, dim=0)
+            split_vs = img_value.split(patch_token_nums, dim=0)
+            for i in range(num_patch):
+                self._kv_cache[block_idx][i] = [
+                    split_ks[i].clone(),
+                    split_vs[i].clone(),
+                ]
+            return key, value
+
+        # ---- Async mode ----
+
+        cur_slot = self.pipeline_state.pipeline_patch_idx
+        if block_idx not in self._kv_cache:
+            self._kv_cache[block_idx] = [None] * num_patch
+
+        # Store fresh K/V in cache (clone for persistence across timesteps)
+        self._kv_cache[block_idx][cur_slot] = [img_key.clone(), img_value.clone()]
+
+        # Build full K/V using pre-allocated buffer + copy_ (avoids torch.cat)
+        total_img = sum(patch_token_nums)
+        full_len = num_txt_tokens + total_img
+
+        if self._full_k_buf is None or self._full_k_buf.shape[0] != full_len or self._full_k_buf.dtype != key.dtype:
+            self._full_k_buf = torch.empty(full_len, *key.shape[1:], dtype=key.dtype, device=key.device)
+            self._full_v_buf = torch.empty(full_len, *value.shape[1:], dtype=value.dtype, device=value.device)
+
+        buf_k = self._full_k_buf
+        buf_v = self._full_v_buf
+
+        # Copy text K/V (fresh, from current patch's computation)
+        buf_k[:num_txt_tokens].copy_(text_key)
+        buf_v[:num_txt_tokens].copy_(text_value)
+
+        # Copy each slot's image K/V into buffer
+        offset = num_txt_tokens
+        for slot in range(num_patch):
+            n = patch_token_nums[slot]
+            if slot == cur_slot:
+                # Fresh from this patch (copy from img_key, already cloned to cache)
+                buf_k[offset : offset + n].copy_(img_key)
+                buf_v[offset : offset + n].copy_(img_value)
+            else:
+                # Stale from cache (previous timestep)
+                cached = self._kv_cache[block_idx][slot]
+                buf_k[offset : offset + n].copy_(cached[0])
+                buf_v[offset : offset + n].copy_(cached[1])
+            offset += n
+
+        return buf_k[:full_len], buf_v[:full_len]
+
+    def clear_kv_cache(self):
+        """Clear stale-KV cache.
+
+        NOTE: stale-KV cache persists ACROSS timesteps by design — that's the
+        whole point of "stale" KV. This method is provided for defensive
+        cleanup only and should NOT be called between timesteps in async mode.
+        """
+        self._kv_cache.clear()
+
+    # ------------------------------------------------------------------
+    # PipeFusion forward
+    # ------------------------------------------------------------------
+
+    def infer(self, block_weights, pre_infer_out):
+        """Run this stage's blocks only.
+
+        Returns ``(hidden_states, encoder_hidden_states, num_txt_tokens)``.
+
+        For non-last stages, streams are ALWAYS split back to (image, text)
+        before returning, so P2P always carries separate streams with
+        consistent shapes.
+        """
+        hidden_states = pre_infer_out.hidden_states
+        encoder_hidden_states = pre_infer_out.encoder_hidden_states
+        timestep = pre_infer_out.timestep
+        image_rotary_emb = pre_infer_out.image_rotary_emb
+
+        # Compute num_txt_tokens
+        if encoder_hidden_states is not None:
+            num_txt_tokens = encoder_hidden_states.shape[0]
+        else:
+            # Streams already concatenated by previous stage — split them
+            txt_ids = pre_infer_out.txt_ids
+            num_txt_tokens = txt_ids.shape[0] if txt_ids is not None else 0
+            if num_txt_tokens > 0:
+                encoder_hidden_states = hidden_states[:num_txt_tokens, ...]
+                hidden_states = hidden_states[num_txt_tokens:, ...]
+
+        image_rotary_emb = self._prepare_image_rotary_emb(image_rotary_emb, num_txt_tokens)
+
+        # Modulation embeddings (computed on every stage)
+        timestep_act = F.silu(timestep)
+        double_stream_mod_img = block_weights.double_stream_modulation_img_linear.apply(timestep_act)
+        double_stream_mod_txt = block_weights.double_stream_modulation_txt_linear.apply(timestep_act)
+        single_stream_mod = block_weights.single_stream_modulation_linear.apply(timestep_act)
+
+        # Double-stream blocks (this stage's subset)
+        for block in block_weights.double_blocks:
+            encoder_hidden_states, hidden_states = self.infer_double_stream_block(
+                block,
+                hidden_states,
+                encoder_hidden_states,
+                double_stream_mod_img,
+                double_stream_mod_txt,
+                image_rotary_emb,
+            )
+
+        # Single-stream blocks: cat [text, image], run, then split back
+        has_single = len(block_weights.single_blocks) > 0
+        if has_single:
+            hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=0)
+
+            for block in block_weights.single_blocks:
+                hidden_states = self.infer_single_stream_block(
+                    block,
+                    hidden_states,
+                    None,
+                    single_stream_mod,
+                    image_rotary_emb,
+                    num_txt_tokens=num_txt_tokens,
+                )
+
+            # Split back to (text, image)
+            encoder_hidden_states = hidden_states[:num_txt_tokens, ...]
+            hidden_states = hidden_states[num_txt_tokens:, ...]
+
+        return hidden_states, encoder_hidden_states, num_txt_tokens
diff --git a/lightx2v/models/networks/flux2/infer/pre_infer.py b/lightx2v/models/networks/flux2/infer/pre_infer.py
index e4fe3b9ef..2bd1bee51 100644
--- a/lightx2v/models/networks/flux2/infer/pre_infer.py
+++ b/lightx2v/models/networks/flux2/infer/pre_infer.py
@@ -76,6 +76,49 @@ def infer(self, weights, hidden_states, encoder_hidden_states, txt_ids=None, img
             image_rotary_emb=image_rotary_emb,
         )
 
+    def infer_partial(self, weights, hidden_states, encoder_hidden_states, txt_ids=None, img_ids=None):
+        """Compute timestep embedding and RoPE only (skip x_embedder / context_embedder).
+
+        Used by non-first pipeline stages that receive already-embedded
+        hidden_states and encoder_hidden_states via P2P.
+        """
+        timesteps_proj = self.scheduler.timesteps_proj
+        timestep_embed = weights.timestep_embedder_linear_1.apply(timesteps_proj)
+        timestep_embed = F.silu(timestep_embed)
+        timestep_embed = weights.timestep_embedder_linear_2.apply(timestep_embed)
+
+        txt_ids_final = txt_ids if txt_ids is not None else getattr(self.scheduler, "txt_ids", None)
+        img_ids_final = img_ids if img_ids is not None else getattr(self.scheduler, "latent_image_ids", None)
+
+        image_rotary_emb = None
+        if img_ids_final is not None and txt_ids_final is not None:
+            if img_ids_final.ndim == 3:
+                img_ids_final = img_ids_final[0]
+            if txt_ids_final.ndim == 3:
+                txt_ids_final = txt_ids_final[0]
+
+            image_rope = self.pos_embed(img_ids_final)
+            text_rope = self.pos_embed(txt_ids_final)
+
+            freqs_cos = torch.cat([text_rope[0], image_rope[0]], dim=0)
+            freqs_sin = torch.cat([text_rope[1], image_rope[1]], dim=0)
+
+        if self.config.get("rope_type", "flashinfer") == "flashinfer":
+            cos_half = freqs_cos[:, ::2].contiguous()
+            sin_half = freqs_sin[:, ::2].contiguous()
+            image_rotary_emb = torch.cat([cos_half, sin_half], dim=-1)
+        else:
+            image_rotary_emb = (freqs_cos, freqs_sin)
+
+        return Flux2PreInferModuleOutput(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            timestep=timestep_embed,
+            txt_ids=txt_ids_final,
+            img_ids=img_ids_final,
+            image_rotary_emb=image_rotary_emb,
+        )
+
 
 class Flux2DevPreInfer(Flux2PreInfer):
     """Pre-processing inference for Flux2 Dev.
diff --git a/lightx2v/models/networks/flux2/infer/transformer_infer.py b/lightx2v/models/networks/flux2/infer/transformer_infer.py
index 5236c2e5a..f0fbb60cf 100644
--- a/lightx2v/models/networks/flux2/infer/transformer_infer.py
+++ b/lightx2v/models/networks/flux2/infer/transformer_infer.py
@@ -33,6 +33,14 @@ def __init__(self, config):
         rope_type = config.get("rope_type", "flashinfer")
         self.apply_rope_func = rope_funcs.get(rope_type, apply_rope_with_torch)
 
+    def _maybe_apply_stale_kv(self, key, value, num_txt_tokens, block_idx):
+        """Hook for stale-KV cache in PipeFusion mode.  No-op in base class.
+
+        Subclasses (PipeFusion) override this to cache image KV across patches
+        while keeping text KV fresh.
+        """
+        return key, value
+
     def set_scheduler(self, scheduler):
         self.scheduler = scheduler
 
@@ -95,8 +103,14 @@ def infer_double_stream_block(
 
         query, key = self.apply_rope_func(query, key, image_rotary_emb)
 
+        # Stale-KV hook (no-op in base class; PipeFusion subclass overrides)
+        num_txt_tokens = encoder_hidden_states.shape[0]
+        key, value = self._maybe_apply_stale_kv(key, value, num_txt_tokens, block_weights.block_idx)
+
         total_len = query.shape[0]
-        cu_seqlens = torch.tensor([0, total_len], dtype=torch.int32)
+        kv_len = key.shape[0]  # may differ from total_len in PipeFusion (stale-KV)
+        cu_seqlens_q = torch.tensor([0, total_len], dtype=torch.int32)
+        cu_seqlens_kv = torch.tensor([0, kv_len], dtype=torch.int32)
 
         model_cls = self.config.get("model_cls", "flux2_klein")
 
@@ -107,7 +121,7 @@ def infer_double_stream_block(
                 k=key,
                 v=value,
                 slice_qkv_len=txt_len,
-                cu_seqlens_qkv=cu_seqlens,
+                cu_seqlens_qkv=cu_seqlens_q,
                 attention_module=block_weights.calculate,
                 seq_p_group=self.seq_p_group,
                 use_fp8_comm=self.seq_p_fp8_comm,
@@ -121,10 +135,10 @@ def infer_double_stream_block(
                 q=query,
                 k=key,
                 v=value,
-                cu_seqlens_q=cu_seqlens,
-                cu_seqlens_kv=cu_seqlens,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_kv=cu_seqlens_kv,
                 max_seqlen_q=total_len,
-                max_seqlen_kv=total_len,
+                max_seqlen_kv=kv_len,
                 model_cls=model_cls,
             )
 
@@ -196,8 +210,13 @@ def infer_single_stream_block(
 
         query, key = self.apply_rope_func(query, key, image_rotary_emb)
 
+        # Stale-KV hook (no-op in base class; PipeFusion subclass overrides)
+        key, value = self._maybe_apply_stale_kv(key, value, num_txt_tokens, block_weights.block_idx)
+
         total_len = query.shape[0]
-        cu_seqlens = torch.tensor([0, total_len], dtype=torch.int32)
+        kv_len = key.shape[0]  # may differ from total_len in PipeFusion (stale-KV)
+        cu_seqlens_q = torch.tensor([0, total_len], dtype=torch.int32)
+        cu_seqlens_kv = torch.tensor([0, kv_len], dtype=torch.int32)
 
         model_cls = self.config.get("model_cls", "flux2_klein")
 
@@ -207,7 +226,7 @@ def infer_single_stream_block(
                 k=key,
                 v=value,
                 slice_qkv_len=num_txt_tokens,
-                cu_seqlens_qkv=cu_seqlens,
+                cu_seqlens_qkv=cu_seqlens_q,
                 attention_module=block_weights.calculate,
                 seq_p_group=self.seq_p_group,
                 use_fp8_comm=self.seq_p_fp8_comm,
@@ -221,10 +240,10 @@ def infer_single_stream_block(
                 q=query,
                 k=key,
                 v=value,
-                cu_seqlens_q=cu_seqlens,
-                cu_seqlens_kv=cu_seqlens,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_kv=cu_seqlens_kv,
                 max_seqlen_q=total_len,
-                max_seqlen_kv=total_len,
+                max_seqlen_kv=kv_len,
                 model_cls=model_cls,
             )
 
diff --git a/lightx2v/models/networks/flux2/model.py b/lightx2v/models/networks/flux2/model.py
index d0c389194..8bdcfe477 100644
--- a/lightx2v/models/networks/flux2/model.py
+++ b/lightx2v/models/networks/flux2/model.py
@@ -26,6 +26,10 @@ def __init__(self, config, model_path, device):
         self.attention_kwargs = {}
         self._init_infer_class()
         self._init_weights()
+        # In PipeFusion mode, weights were loaded to CPU to avoid OOM;
+        # move only this stage's subset to GPU.
+        if self.config.get("pipefusion_parallel", False):
+            self.to_cuda()
         self._init_infer()
 
     def _init_infer(self):
@@ -105,7 +109,13 @@ class Flux2KleinTransformerModel(_Flux2TransformerModelBase):
 
     def _init_infer_class(self):
         feature_caching = self.config.get("feature_caching", "NoCaching")
-        if feature_caching in ("NoCaching", "None"):
+        if self.config.get("pipefusion_parallel", False):
+            from lightx2v.models.networks.flux2.infer.pipefusion.transformer_infer import (
+                Flux2PipeFusionTransformerInfer,
+            )
+
+            self.transformer_infer_class = Flux2PipeFusionTransformerInfer
+        elif feature_caching in ("NoCaching", "None"):
             if self.cpu_offload and self.offload_granularity == "block":
                 self.transformer_infer_class = Flux2OffloadTransformerInfer
             else:
diff --git a/lightx2v/models/networks/flux2/weights/transformer_weights.py b/lightx2v/models/networks/flux2/weights/transformer_weights.py
index 2777c357f..324b4ccb9 100644
--- a/lightx2v/models/networks/flux2/weights/transformer_weights.py
+++ b/lightx2v/models/networks/flux2/weights/transformer_weights.py
@@ -271,8 +271,50 @@ def __init__(self, config):
 
         inner_dim = config.get("num_attention_heads", 24) * config.get("attention_head_dim", 64)
 
-        self.double_blocks = WeightModuleList([Flux2DoubleBlockWeights(config, i) for i in range(self.num_layers)])
-        self.single_blocks = WeightModuleList([Flux2SingleBlockWeights(config, i) for i in range(self.num_single_layers)])
+        # -- Pipeline-parallel block splitting --------------------------------
+        pp_size = config.get("pipefusion_parallel", False)
+        if pp_size:
+            from lightx2v.common.distributed import (
+                get_pipeline_parallel_rank,
+                get_pipeline_parallel_world_size,
+            )
+
+            pp_rank = get_pipeline_parallel_rank()
+            pp_world_size = get_pipeline_parallel_world_size()
+        else:
+            pp_rank = 0
+            pp_world_size = 1
+
+        if pp_world_size > 1:
+            # Split double_blocks + single_blocks across pipeline stages.
+            # Blocks are assigned contiguously: stage 0 gets the first chunk,
+            # stage 1 the next, etc.  A stage may span the double→single
+            # boundary (it will then have both types).
+            total_blocks = self.num_layers + self.num_single_layers
+            blocks_per_stage = (total_blocks + pp_world_size - 1) // pp_world_size
+            stage_start = pp_rank * blocks_per_stage
+            stage_end = min((pp_rank + 1) * blocks_per_stage, total_blocks)
+
+            double_start = min(stage_start, self.num_layers)
+            double_end = min(stage_end, self.num_layers)
+            single_start = max(0, stage_start - self.num_layers)
+            single_end = max(0, stage_end - self.num_layers)
+
+            self.double_blocks = WeightModuleList([Flux2DoubleBlockWeights(config, i) for i in range(double_start, double_end)])
+            self.single_blocks = WeightModuleList([Flux2SingleBlockWeights(config, i) for i in range(single_start, single_end)])
+            # Track whether this stage crosses the double→single boundary
+            self._has_double = double_end > double_start
+            self._has_single = single_end > single_start
+            self._stage_start = stage_start
+            self._stage_end = stage_end
+        else:
+            self.double_blocks = WeightModuleList([Flux2DoubleBlockWeights(config, i) for i in range(self.num_layers)])
+            self.single_blocks = WeightModuleList([Flux2SingleBlockWeights(config, i) for i in range(self.num_single_layers)])
+            self._has_double = True
+            self._has_single = True
+            self._stage_start = 0
+            self._stage_end = self.num_layers + self.num_single_layers
+
         self.register_offload_buffers(config)
 
         self.add_module("double_blocks", self.double_blocks)
diff --git a/lightx2v/models/runners/flux2/flux2_runner.py b/lightx2v/models/runners/flux2/flux2_runner.py
index 33d84d72a..197238ad6 100644
--- a/lightx2v/models/runners/flux2/flux2_runner.py
+++ b/lightx2v/models/runners/flux2/flux2_runner.py
@@ -243,6 +243,12 @@ def _run_dit_local_i2i(self, total_steps=None):
         return latents, generator
 
     def run(self, total_steps=None):
+        if self.config.get("pipefusion_parallel", False):
+            return self._run_pipefusion(total_steps)
+        return self._run_sequential(total_steps)
+
+    def _run_sequential(self, total_steps=None):
+        """Existing synchronous denoising loop (single-GPU or non-PipeFusion)."""
         if total_steps is None:
             total_steps = self.model.scheduler.infer_steps
         for step_index in range(total_steps):
@@ -262,6 +268,78 @@ def run(self, total_steps=None):
 
         return self.model.scheduler.latents, self.model.scheduler.generator
 
+    def _run_pipefusion(self, total_steps=None):
+        """PipeFusion denoising loop: pipeline driver controls all timesteps."""
+        from lightx2v.common.distributed import (
+            get_pipeline_runtime_state,
+            is_pipeline_last_stage,
+        )
+
+        if total_steps is None:
+            total_steps = self.model.scheduler.infer_steps
+
+        # Initialize pipeline runtime state with image dimensions
+        pipeline_state = get_pipeline_runtime_state()
+        height = self.input_info.latent_shape[1]  # packed_h * packed_w tokens
+        # Reconstruct actual height/width from latent_image_ids
+        latent_image_ids = self.model.scheduler.latent_image_ids
+        if latent_image_ids is not None:
+            packed_h = int((latent_image_ids[0, :, 1].max() + 1).item()) if latent_image_ids.ndim == 3 else int((latent_image_ids[:, 1].max() + 1).item())
+            packed_w = int((latent_image_ids[0, :, 2].max() + 1).item()) if latent_image_ids.ndim == 3 else int((latent_image_ids[:, 2].max() + 1).item())
+            vae_scale_factor = self.config.get("vae_scale_factor", 16)
+            actual_height = packed_h * vae_scale_factor * 2
+            actual_width = packed_w * vae_scale_factor * 2
+        else:
+            actual_height = actual_width = 1024
+
+        num_pipeline_patch = self.config.get("parallel", {}).get("num_pipeline_patch", 4)
+        warmup_steps = self.config.get("parallel", {}).get("pipeline_warmup_steps", 1)
+
+        pipeline_state.set_input_parameters(
+            height=actual_height,
+            width=actual_width,
+            batch_size=1,
+            num_pipeline_patch=num_pipeline_patch,
+            warmup_steps=warmup_steps,
+            vae_scale_factor=self.config.get("vae_scale_factor", 16),
+        )
+
+        # Prepare inputs
+        latents = self.model.scheduler.latents
+        text_encoder_output = self.inputs["text_encoder_output"]
+        prompt_embeds = text_encoder_output["prompt_embeds"]
+        text_ids = text_encoder_output.get("text_ids")
+        latent_image_ids = self.model.scheduler.latent_image_ids
+
+        do_cfg = self.config.get("enable_cfg", True) and self.config.get("sample_guide_scale", 1.0) > 1.0
+        negative_prompt_embeds = text_encoder_output.get("negative_prompt_embeds") if do_cfg else None
+        negative_text_ids = text_encoder_output.get("negative_text_ids") if do_cfg else None
+
+        timesteps = self.model.scheduler.timesteps
+
+        # Run pipeline
+        from lightx2v.models.networks.flux2.infer.pipefusion.pipeline_driver import (
+            Flux2PipelineDriver,
+        )
+
+        driver = Flux2PipelineDriver(self.model, self.config)
+        latents = driver.run_pipeline(
+            latents=latents,
+            prompt_embeds=prompt_embeds,
+            text_ids=text_ids,
+            latent_image_ids=latent_image_ids,
+            timesteps=timesteps,
+            scheduler=self.model.scheduler,
+            do_cfg=do_cfg,
+            negative_prompt_embeds=negative_prompt_embeds,
+            negative_text_ids=negative_text_ids,
+        )
+
+        if latents is not None and is_pipeline_last_stage():
+            self.model.scheduler.latents = latents
+
+        return self.model.scheduler.latents, self.model.scheduler.generator
+
     def get_custom_shape(self):
         default_aspect_ratios = {
             "16:9": [1344, 768],
@@ -367,13 +445,27 @@ def run_pipeline(self, input_info):
         self.set_img_shapes()
 
         latents, generator = self.run_dit()
-        images = self.run_vae_decoder(latents)
+
+        # In PipeFusion mode, only the last stage has final latents
+        if self.config.get("pipefusion_parallel", False):
+            from lightx2v.common.distributed import is_pipeline_last_stage
+
+            if is_pipeline_last_stage():
+                images = self.run_vae_decoder(latents)
+            else:
+                images = None
+        else:
+            images = self.run_vae_decoder(latents)
         self.end_run()
 
-        if not input_info.return_result_tensor and is_main_process():
-            image = images[0]
-            image.save(input_info.save_result_path)
-            logger.info(f"Image saved: {input_info.save_result_path}")
+        # Save image: in PipeFusion mode, last stage has the image;
+        # in normal mode, main process (rank 0) has it.
+        if not input_info.return_result_tensor:
+            should_save = is_pipeline_last_stage() if self.config.get("pipefusion_parallel", False) else is_main_process()
+            if should_save and images is not None:
+                image = images[0]
+                image.save(input_info.save_result_path)
+                logger.info(f"Image saved: {input_info.save_result_path}")
 
         del latents, generator
         torch_device_module.empty_cache()
diff --git a/lightx2v/models/schedulers/flux2/scheduler.py b/lightx2v/models/schedulers/flux2/scheduler.py
index 32e007cde..432bf4321 100755
--- a/lightx2v/models/schedulers/flux2/scheduler.py
+++ b/lightx2v/models/schedulers/flux2/scheduler.py
@@ -146,6 +146,15 @@ def step_post(self):
         )
         self.latents = latents
 
+    def step_post_patch(self, noise_pred, latents, t):
+        """Patch-level scheduler step for async PipeFusion mode.
+
+        Unlike ``step_post``, this operates on a single patch's latents and
+        noise_pred, and does not apply FLS enhancement (which requires the
+        full latent).
+        """
+        return self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
     def _encode_image(self, image):
         image = image.to(device=AI_DEVICE, dtype=GET_DTYPE())
         encoder_output = self.vae.encode_vae_image(image)
diff --git a/lightx2v/pipeline.py b/lightx2v/pipeline.py
index 351bb31dd..8365c200b 100755
--- a/lightx2v/pipeline.py
+++ b/lightx2v/pipeline.py
@@ -422,11 +422,14 @@ def enable_cache(
             self.magcache_retention_ratio = magcache_retention_ratio
             self.magcache_ratios = magcache_ratios
 
-    def enable_parallel(self, cfg_p_size=1, seq_p_size=1, seq_p_attn_type="ulysses"):
+    def enable_parallel(self, cfg_p_size=1, seq_p_size=1, seq_p_attn_type="ulysses", pp_size=1, num_pipeline_patch=4, pipeline_warmup_steps=1):
         self.parallel = {
             "cfg_p_size": cfg_p_size,
             "seq_p_size": seq_p_size,
             "seq_p_attn_type": seq_p_attn_type,
+            "pp_size": pp_size,
+            "num_pipeline_patch": num_pipeline_patch,
+            "pipeline_warmup_steps": pipeline_warmup_steps,
         }
 
     @torch.no_grad()
diff --git a/lightx2v/utils/set_config.py b/lightx2v/utils/set_config.py
index 97155fa10..e10aef058 100755
--- a/lightx2v/utils/set_config.py
+++ b/lightx2v/utils/set_config.py
@@ -28,6 +28,7 @@ def get_default_config():
         "parallel": False,
         "seq_parallel": False,
         "cfg_parallel": False,
+        "pipefusion_parallel": False,
         "enable_cfg": False,
         "use_image_encoder": True,
     }
@@ -280,12 +281,16 @@ def set_parallel_config(config):
             config["tensor_parallel"] = True
             config["seq_parallel"] = False
             config["cfg_parallel"] = False
+            config["pipefusion_parallel"] = False
         else:
-            # Original 2D mesh for cfg_p and seq_p
+            # Multi-dimensional mesh: (cfg_p, pp, seq_p)
             cfg_p_size = config["parallel"].get("cfg_p_size", 1)
+            pp_size = config["parallel"].get("pp_size", 1)
             seq_p_size = config["parallel"].get("seq_p_size", 1)
-            assert cfg_p_size * seq_p_size == dist.get_world_size(), f"cfg_p_size ({cfg_p_size}) * seq_p_size ({seq_p_size}) must be equal to world_size ({dist.get_world_size()})"
-            config["device_mesh"] = init_device_mesh(AI_DEVICE, (cfg_p_size, seq_p_size), mesh_dim_names=("cfg_p", "seq_p"))
+            assert cfg_p_size * pp_size * seq_p_size == dist.get_world_size(), (
+                f"cfg_p_size ({cfg_p_size}) * pp_size ({pp_size}) * seq_p_size ({seq_p_size}) must be equal to world_size ({dist.get_world_size()})"
+            )
+            config["device_mesh"] = init_device_mesh(AI_DEVICE, (cfg_p_size, pp_size, seq_p_size), mesh_dim_names=("cfg_p", "pp", "seq_p"))
             config["tensor_parallel"] = False
 
             if config["parallel"] and config["parallel"].get("seq_p_size", False) and config["parallel"]["seq_p_size"] > 1:
@@ -294,9 +299,18 @@ def set_parallel_config(config):
             if config.get("enable_cfg", False) and config["parallel"] and config["parallel"].get("cfg_p_size", False) and config["parallel"]["cfg_p_size"] > 1:
                 config["cfg_parallel"] = True
 
+            config["pipefusion_parallel"] = pp_size > 1
+            if pp_size > 1:
+                from lightx2v.common.distributed import init_pipeline_parallel_state
+
+                pp_group = config["device_mesh"].get_group(mesh_dim="pp")
+                init_pipeline_parallel_state(pp_group)
+
         # warmup dist
         _a = torch.zeros([1]).to(f"{AI_DEVICE}:{dist.get_rank()}")
         dist.all_reduce(_a)
+    else:
+        config["pipefusion_parallel"] = False
 
 
 def print_config(config):

From 0bb43332bcb0d685f4fb37e7290070e3f43738c6 Mon Sep 17 00:00:00 2001
From: Qin-sx <67671068+Qin-sx@users.noreply.github.com>
Date: Mon, 20 Jul 2026 21:53:50 +0800
Subject: [PATCH 02/12] Update
 lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 .../networks/flux2/infer/pipefusion/transformer_infer.py    | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py b/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
index 5c1bbb684..250362b5a 100644
--- a/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
+++ b/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
@@ -32,10 +32,10 @@ def __init__(self, config):
         # Stale-KV cache: block_idx -> [ [k, v] per patch slot ]
         self._kv_cache: dict = {}
 
-        # Pre-allocated full K/V buffers (lazily created on first async use).
+        # Pre-allocated full K/V buffers per block (lazily created on first async use).
         # Avoids repeated torch.cat allocations per timestep.
-        self._full_k_buf = None
-        self._full_v_buf = None
+        self._full_k_bufs: dict = {}
+        self._full_v_bufs: dict = {}
 
     # ------------------------------------------------------------------
     # Stale-KV hook (overrides base class no-op)

From cfdce1410bca70f58a87e1cfa53bfb75f38d4e53 Mon Sep 17 00:00:00 2001
From: Qin-sx <67671068+Qin-sx@users.noreply.github.com>
Date: Mon, 20 Jul 2026 21:59:52 +0800
Subject: [PATCH 03/12] Update
 lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 .../flux2/infer/pipefusion/transformer_infer.py        | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py b/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
index 250362b5a..464c8dcca 100644
--- a/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
+++ b/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
@@ -93,12 +93,12 @@ def _maybe_apply_stale_kv(self, key, value, num_txt_tokens, block_idx):
         total_img = sum(patch_token_nums)
         full_len = num_txt_tokens + total_img
 
-        if self._full_k_buf is None or self._full_k_buf.shape[0] != full_len or self._full_k_buf.dtype != key.dtype:
-            self._full_k_buf = torch.empty(full_len, *key.shape[1:], dtype=key.dtype, device=key.device)
-            self._full_v_buf = torch.empty(full_len, *value.shape[1:], dtype=value.dtype, device=value.device)
+        if block_idx not in self._full_k_bufs or self._full_k_bufs[block_idx].shape[0] != full_len or self._full_k_bufs[block_idx].dtype != key.dtype:
+            self._full_k_bufs[block_idx] = torch.empty(full_len, *key.shape[1:], dtype=key.dtype, device=key.device)
+            self._full_v_bufs[block_idx] = torch.empty(full_len, *value.shape[1:], dtype=value.dtype, device=value.device)
 
-        buf_k = self._full_k_buf
-        buf_v = self._full_v_buf
+        buf_k = self._full_k_bufs[block_idx]
+        buf_v = self._full_v_bufs[block_idx]
 
         # Copy text K/V (fresh, from current patch's computation)
         buf_k[:num_txt_tokens].copy_(text_key)

From 8a7e8540948a60803f2bf2fcf1bebfb28d3e8679 Mon Sep 17 00:00:00 2001
From: Qin-sx <67671068+Qin-sx@users.noreply.github.com>
Date: Mon, 20 Jul 2026 22:06:34 +0800
Subject: [PATCH 04/12] Update
 lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 .../models/networks/flux2/infer/pipefusion/pipeline_driver.py   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py b/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
index d60d47220..e662e2891 100644
--- a/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
+++ b/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
@@ -342,7 +342,7 @@ def _async_pipeline(self, latents, prompt_embeds, text_ids, latent_image_ids, ti
                             self.pp_comm.recv_next()
 
                 # ---- 5. Wait for old isends (limit pending to prevent GC issues) ----
-                while len(pending_isends) > 4:
+                while len(pending_isends) > num_patch * 2:
                     old_req, _ = pending_isends.pop(0)
                     old_req.wait()
 

From e05d42189c367e2afea2aed6e4f3541c0b8083c4 Mon Sep 17 00:00:00 2001
From: Qin-sx <67671068+Qin-sx@users.noreply.github.com>
Date: Mon, 20 Jul 2026 22:07:33 +0800
Subject: [PATCH 05/12] Update
 lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 .../models/networks/flux2/infer/pipefusion/transformer_infer.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py b/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
index 464c8dcca..57604a63b 100644
--- a/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
+++ b/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
@@ -129,6 +129,8 @@ def clear_kv_cache(self):
         cleanup only and should NOT be called between timesteps in async mode.
         """
         self._kv_cache.clear()
+        self._full_k_bufs.clear()
+        self._full_v_bufs.clear()
 
     # ------------------------------------------------------------------
     # PipeFusion forward

From a5087e056e24252541c987ced95bc13272cd648c Mon Sep 17 00:00:00 2001
From: Qin-sx <67671068+Qin-sx@users.noreply.github.com>
Date: Mon, 20 Jul 2026 22:09:52 +0800
Subject: [PATCH 06/12] Update
 lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 .../models/networks/flux2/infer/pipefusion/pipeline_driver.py   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py b/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
index e662e2891..eb9224cbf 100644
--- a/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
+++ b/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
@@ -315,7 +315,7 @@ def _async_pipeline(self, latents, prompt_embeds, text_ids, latent_image_ids, ti
                 if self._is_last:
                     noise_pred = result
                     scheduler.scheduler._step_index = i + self.state.warmup_steps
-                    patch_latents[patch_idx] = scheduler.scheduler.step(noise_pred, t, last_patch_latents[patch_idx], return_dict=False)[0]
+                    patch_latents[patch_idx] = scheduler.step_post_patch(noise_pred, last_patch_latents[patch_idx], t)
                     if i != total_steps - 1:
                         req = self.pp_comm.pipeline_isend(patch_latents[patch_idx], name="latent", segment_idx=patch_idx)
                         pending_isends.append((req, patch_latents[patch_idx]))

From ab9baaf227ba845fdb7b480e67299d8aac83b895 Mon Sep 17 00:00:00 2001
From: Qin-sx <67671068+Qin-sx@users.noreply.github.com>
Date: Mon, 20 Jul 2026 22:22:07 +0800
Subject: [PATCH 07/12] Update
 lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 .../models/networks/flux2/infer/pipefusion/pipeline_driver.py   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py b/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
index eb9224cbf..a512db1d3 100644
--- a/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
+++ b/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
@@ -223,7 +223,7 @@ def _async_pipeline(self, latents, prompt_embeds, text_ids, latent_image_ids, ti
         inner_dim = self.config.get("num_attention_heads", 24) * self.config.get("attention_head_dim", 64)
         # Raw latent channels (before x_embedder): rank=0 recv from last stage
         # gets raw latents [1, L, C_in]; other stages recv embedded [L, D]
-        raw_channels = self.config.get("transformer_in_channels", self.config.get("in_channels", 128))
+        raw_channels = getattr(self.model, "in_channels", self.config.get("transformer_in_channels", self.config.get("in_channels", 128)))
 
         # Split latents into patches (dim=1 for [B, L, C])
         if self._is_first or self._is_last:

From 3e6f7c5847bbd8422b072db223c3422c38d4ad5b Mon Sep 17 00:00:00 2001
From: Qin-sx <Qin_sx@outlook.com>
Date: Mon, 20 Jul 2026 22:27:58 +0800
Subject: [PATCH 08/12] update lightx2v/models/runners/flux2/flux2_runner.py

	modified:   lightx2v/models/runners/flux2/flux2_runner.py
---
 lightx2v/models/runners/flux2/flux2_runner.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/lightx2v/models/runners/flux2/flux2_runner.py b/lightx2v/models/runners/flux2/flux2_runner.py
index 197238ad6..dbbb0525c 100644
--- a/lightx2v/models/runners/flux2/flux2_runner.py
+++ b/lightx2v/models/runners/flux2/flux2_runner.py
@@ -283,12 +283,8 @@ def _run_pipefusion(self, total_steps=None):
         height = self.input_info.latent_shape[1]  # packed_h * packed_w tokens
         # Reconstruct actual height/width from latent_image_ids
         latent_image_ids = self.model.scheduler.latent_image_ids
-        if latent_image_ids is not None:
-            packed_h = int((latent_image_ids[0, :, 1].max() + 1).item()) if latent_image_ids.ndim == 3 else int((latent_image_ids[:, 1].max() + 1).item())
-            packed_w = int((latent_image_ids[0, :, 2].max() + 1).item()) if latent_image_ids.ndim == 3 else int((latent_image_ids[:, 2].max() + 1).item())
-            vae_scale_factor = self.config.get("vae_scale_factor", 16)
-            actual_height = packed_h * vae_scale_factor * 2
-            actual_width = packed_w * vae_scale_factor * 2
+        if self.input_info.target_shape is not None:
+            actual_height, actual_width = self.input_info.target_shape
         else:
             actual_height = actual_width = 1024
 

From 448b144d7771b81a05b220c3ec91619ad9e0a2d6 Mon Sep 17 00:00:00 2001
From: Qin-sx <Qin_sx@outlook.com>
Date: Mon, 20 Jul 2026 23:10:13 +0800
Subject: [PATCH 09/12] update
 lightx2v/models/networks/flux2/weights/transformer_weights.py 	modified:  
 lightx2v/models/networks/flux2/weights/transformer_weights.py

---
 lightx2v/models/networks/flux2/weights/transformer_weights.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/lightx2v/models/networks/flux2/weights/transformer_weights.py b/lightx2v/models/networks/flux2/weights/transformer_weights.py
index e7170d077..8221a82ce 100644
--- a/lightx2v/models/networks/flux2/weights/transformer_weights.py
+++ b/lightx2v/models/networks/flux2/weights/transformer_weights.py
@@ -154,8 +154,6 @@ def __init__(self, config):
         self.num_single_layers = config.get("num_single_layers", 20)
         self.mm_type = config.get("dit_quant_scheme", "Default")
 
-        inner_dim = config.get("num_attention_heads", 24) * config.get("attention_head_dim", 64)
-
         # -- Pipeline-parallel block splitting --------------------------------
         pp_size = config.get("pipefusion_parallel", False)
         if pp_size:

From 7fa04ed707756c22b2e89a6d9b58824644e4e75c Mon Sep 17 00:00:00 2001
From: Qin-sx <Qin_sx@outlook.com>
Date: Fri, 24 Jul 2026 09:01:54 +0800
Subject: [PATCH 10/12] update for PRs

	modified:   lightx2v/common/distributed/pipeline_state.py
	modified:   lightx2v/common/ops/attn/flash_attn.py
	modified:   lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
	modified:   lightx2v/models/networks/flux2/infer/pre_infer.py
	modified:   lightx2v/models/networks/flux2/model.py
	modified:   lightx2v/models/networks/flux2/weights/transformer_weights.py
	modified:   lightx2v/models/runners/flux2/flux2_runner.py
---
 lightx2v/common/distributed/pipeline_state.py | 20 +++++++++-----
 lightx2v/common/ops/attn/flash_attn.py        |  8 +++---
 .../infer/pipefusion/transformer_infer.py     |  5 ++--
 .../models/networks/flux2/infer/pre_infer.py  | 26 +++++--------------
 lightx2v/models/networks/flux2/model.py       |  3 ++-
 .../flux2/weights/transformer_weights.py      |  7 ++---
 lightx2v/models/runners/flux2/flux2_runner.py |  1 +
 7 files changed, 35 insertions(+), 35 deletions(-)

diff --git a/lightx2v/common/distributed/pipeline_state.py b/lightx2v/common/distributed/pipeline_state.py
index 9838f7ce7..e84ba4b85 100644
--- a/lightx2v/common/distributed/pipeline_state.py
+++ b/lightx2v/common/distributed/pipeline_state.py
@@ -101,6 +101,7 @@ def set_input_parameters(
         warmup_steps: int = 1,
         vae_scale_factor: int = 16,
         patch_size: int = 1,
+        total_tokens: Optional[int] = None,
     ):
         self.height = height
         self.width = width
@@ -111,15 +112,20 @@ def set_input_parameters(
         if num_pipeline_patch is not None:
             self.num_pipeline_patch = num_pipeline_patch
 
-        # Compute packed dimensions (matching Flux2Runner.set_target_shape)
-        multiple_of = vae_scale_factor * 2
-        self.packed_h = height // multiple_of
-        self.packed_w = width // multiple_of
-        total_tokens = self.packed_h * self.packed_w
+        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 = total_tokens // self.num_pipeline_patch
-        remainder = total_tokens % self.num_pipeline_patch
+        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
diff --git a/lightx2v/common/ops/attn/flash_attn.py b/lightx2v/common/ops/attn/flash_attn.py
index 8ed341848..bde3945aa 100755
--- a/lightx2v/common/ops/attn/flash_attn.py
+++ b/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
 
 
diff --git a/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py b/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
index 57604a63b..faad09e18 100644
--- a/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
+++ b/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
@@ -149,6 +149,7 @@ def infer(self, block_weights, pre_infer_out):
         encoder_hidden_states = pre_infer_out.encoder_hidden_states
         timestep = pre_infer_out.timestep
         image_rotary_emb = pre_infer_out.image_rotary_emb
+        image_rotary_positions = pre_infer_out.image_rotary_positions
 
         # Compute num_txt_tokens
         if encoder_hidden_states is not None:
@@ -161,8 +162,6 @@ def infer(self, block_weights, pre_infer_out):
                 encoder_hidden_states = hidden_states[:num_txt_tokens, ...]
                 hidden_states = hidden_states[num_txt_tokens:, ...]
 
-        image_rotary_emb = self._prepare_image_rotary_emb(image_rotary_emb, num_txt_tokens)
-
         # Modulation embeddings (computed on every stage)
         timestep_act = F.silu(timestep)
         double_stream_mod_img = block_weights.double_stream_modulation_img_linear.apply(timestep_act)
@@ -178,6 +177,7 @@ def infer(self, block_weights, pre_infer_out):
                 double_stream_mod_img,
                 double_stream_mod_txt,
                 image_rotary_emb,
+                image_rotary_positions,
             )
 
         # Single-stream blocks: cat [text, image], run, then split back
@@ -192,6 +192,7 @@ def infer(self, block_weights, pre_infer_out):
                     None,
                     single_stream_mod,
                     image_rotary_emb,
+                    image_rotary_positions,
                     num_txt_tokens=num_txt_tokens,
                 )
 
diff --git a/lightx2v/models/networks/flux2/infer/pre_infer.py b/lightx2v/models/networks/flux2/infer/pre_infer.py
index 6c31a685f..a55cb71b8 100644
--- a/lightx2v/models/networks/flux2/infer/pre_infer.py
+++ b/lightx2v/models/networks/flux2/infer/pre_infer.py
@@ -149,25 +149,12 @@ def infer_partial(self, weights, hidden_states, encoder_hidden_states, txt_ids=N
         txt_ids_final = txt_ids if txt_ids is not None else getattr(self.scheduler, "txt_ids", None)
         img_ids_final = img_ids if img_ids is not None else getattr(self.scheduler, "latent_image_ids", None)
 
-        image_rotary_emb = None
-        if img_ids_final is not None and txt_ids_final is not None:
-            if img_ids_final.ndim == 3:
-                img_ids_final = img_ids_final[0]
-            if txt_ids_final.ndim == 3:
-                txt_ids_final = txt_ids_final[0]
-
-            image_rope = self.pos_embed(img_ids_final)
-            text_rope = self.pos_embed(txt_ids_final)
-
-            freqs_cos = torch.cat([text_rope[0], image_rope[0]], dim=0)
-            freqs_sin = torch.cat([text_rope[1], image_rope[1]], dim=0)
-
-        if self.config.get("rope_type", "flashinfer") == "flashinfer":
-            cos_half = freqs_cos[:, ::2].contiguous()
-            sin_half = freqs_sin[:, ::2].contiguous()
-            image_rotary_emb = torch.cat([cos_half, sin_half], dim=-1)
-        else:
-            image_rotary_emb = (freqs_cos, freqs_sin)
+        num_txt_tokens = encoder_hidden_states.shape[0] if encoder_hidden_states is not None else 0
+        image_rotary_emb, image_rotary_positions = self.get_rope_cache(txt_ids_final, img_ids_final, num_txt_tokens)
+        if img_ids_final is not None and img_ids_final.ndim == 3:
+            img_ids_final = img_ids_final[0]
+        if txt_ids_final is not None and txt_ids_final.ndim == 3:
+            txt_ids_final = txt_ids_final[0]
 
         return Flux2PreInferModuleOutput(
             hidden_states=hidden_states,
@@ -176,6 +163,7 @@ def infer_partial(self, weights, hidden_states, encoder_hidden_states, txt_ids=N
             txt_ids=txt_ids_final,
             img_ids=img_ids_final,
             image_rotary_emb=image_rotary_emb,
+            image_rotary_positions=image_rotary_positions,
         )
 
 
diff --git a/lightx2v/models/networks/flux2/model.py b/lightx2v/models/networks/flux2/model.py
index 90469dbf9..88d18e568 100644
--- a/lightx2v/models/networks/flux2/model.py
+++ b/lightx2v/models/networks/flux2/model.py
@@ -276,7 +276,8 @@ def _init_infer(self):
         self.transformer_infer = self.transformer_infer_class(self.config)
         self.pre_infer = self.pre_infer_class(self.config)
         self.post_infer = self.post_infer_class(self.config)
-        self.pre_infer.set_rope(self.transformer_weights.double_blocks[0].rope)
+        blocks = self.transformer_weights.double_blocks or self.transformer_weights.single_blocks
+        self.pre_infer.set_rope(blocks[0].rope)
         if hasattr(self.transformer_infer, "offload_manager_double") and hasattr(self.transformer_infer, "offload_manager_single"):
             self._init_offload_manager()
 
diff --git a/lightx2v/models/networks/flux2/weights/transformer_weights.py b/lightx2v/models/networks/flux2/weights/transformer_weights.py
index 90fe0e2e8..d9e15b298 100644
--- a/lightx2v/models/networks/flux2/weights/transformer_weights.py
+++ b/lightx2v/models/networks/flux2/weights/transformer_weights.py
@@ -183,9 +183,10 @@ def __init__(self, config):
             # stage 1 the next, etc.  A stage may span the double→single
             # boundary (it will then have both types).
             total_blocks = self.num_layers + self.num_single_layers
-            blocks_per_stage = (total_blocks + pp_world_size - 1) // pp_world_size
-            stage_start = pp_rank * blocks_per_stage
-            stage_end = min((pp_rank + 1) * blocks_per_stage, total_blocks)
+            base = total_blocks // pp_world_size
+            remainder = total_blocks % pp_world_size
+            stage_start = pp_rank * base + min(pp_rank, remainder)
+            stage_end = stage_start + base + (1 if pp_rank < remainder else 0)
 
             double_start = min(stage_start, self.num_layers)
             double_end = min(stage_end, self.num_layers)
diff --git a/lightx2v/models/runners/flux2/flux2_runner.py b/lightx2v/models/runners/flux2/flux2_runner.py
index db0cdc6f1..a20de2f3c 100644
--- a/lightx2v/models/runners/flux2/flux2_runner.py
+++ b/lightx2v/models/runners/flux2/flux2_runner.py
@@ -429,6 +429,7 @@ def _run_pipefusion(self, total_steps=None):
             num_pipeline_patch=num_pipeline_patch,
             warmup_steps=warmup_steps,
             vae_scale_factor=self.config.get("vae_scale_factor", 16),
+            total_tokens=self.input_info.latent_shape[1],
         )
 
         # Prepare inputs

From 8cadcfbe2b7ae5f8dd3e5016dce1090b6a30f28e Mon Sep 17 00:00:00 2001
From: Qin-sx <Qin_sx@outlook.com>
Date: Sat, 25 Jul 2026 22:03:48 +0800
Subject: [PATCH 11/12] Offload transformer weights and clear KV cache before
 VAE decode to avoid OOM

	modified:   lightx2v/models/runners/flux2/flux2_runner.py
---
 lightx2v/models/runners/flux2/flux2_runner.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/lightx2v/models/runners/flux2/flux2_runner.py b/lightx2v/models/runners/flux2/flux2_runner.py
index a20de2f3c..65369e1c9 100644
--- a/lightx2v/models/runners/flux2/flux2_runner.py
+++ b/lightx2v/models/runners/flux2/flux2_runner.py
@@ -591,6 +591,12 @@ def run_pipeline(self, input_info):
             from lightx2v.common.distributed import is_pipeline_last_stage
 
             if is_pipeline_last_stage():
+                # Offload transformer weights and clear KV cache before VAE decode to avoid OOM
+                self.model.transformer_weights.to_cpu()
+                if hasattr(self.model.transformer_infer, "clear_kv_cache"):
+                    self.model.transformer_infer.clear_kv_cache()
+                torch_device_module.empty_cache()
+                gc.collect()
                 images = self.run_vae_decoder(latents)
             else:
                 images = None

From 51b6fd59c7553d28feb2a51a191e78c64f4e71b8 Mon Sep 17 00:00:00 2001
From: Qin-sx <Qin_sx@outlook.com>
Date: Sat, 1 Aug 2026 19:01:45 +0800
Subject: [PATCH 12/12] update lightx2v/utils/set_config.py

	modified:   lightx2v/utils/set_config.py
---
 lightx2v/utils/set_config.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/lightx2v/utils/set_config.py b/lightx2v/utils/set_config.py
index 2866369bd..516a527ed 100755
--- a/lightx2v/utils/set_config.py
+++ b/lightx2v/utils/set_config.py
@@ -397,11 +397,12 @@ def set_parallel_config(config):
         tensor_p_size = int(config["parallel"].get("tensor_p_size", 1))
         cfg_p_size = int(config["parallel"].get("cfg_p_size", 1))
         seq_p_size = int(config["parallel"].get("seq_p_size", 1))
+        pp_size = int(config["parallel"].get("pp_size", 1))
         world_size = dist.get_world_size()
-        expected_world_size = tensor_p_size * cfg_p_size * seq_p_size
+        expected_world_size = tensor_p_size * cfg_p_size * seq_p_size * pp_size
         if expected_world_size != world_size:
             raise ValueError(
-                f"Parallel sizes must match the distributed world size: tensor_p_size ({tensor_p_size}) * cfg_p_size ({cfg_p_size}) * seq_p_size ({seq_p_size}) != world_size ({world_size})."
+                f"Parallel sizes must match the distributed world size: tensor_p_size ({tensor_p_size}) * cfg_p_size ({cfg_p_size}) * seq_p_size ({seq_p_size}) * pp_size ({pp_size}) != world_size ({world_size})."
             )
 
         if tensor_p_size > 1: