From 671fb03fbdf3125ed428a438f2470f8b799e0c86 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/11] 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                  |  23 +-
 15 files changed, 1130 insertions(+), 24 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 35a8ceaf7..8dd4fb469 100644
--- a/lightx2v/models/networks/base_model.py
+++ b/lightx2v/models/networks/base_model.py
@@ -403,7 +403,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 651c6a95f..6c31a685f 100644
--- a/lightx2v/models/networks/flux2/infer/pre_infer.py
+++ b/lightx2v/models/networks/flux2/infer/pre_infer.py
@@ -135,6 +135,49 @@ def infer(self, weights, hidden_states, encoder_hidden_states, txt_ids=None, img
             image_rotary_positions=image_rotary_positions,
         )
 
+    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 ba8a754eb..a37695a17 100644
--- a/lightx2v/models/networks/flux2/infer/transformer_infer.py
+++ b/lightx2v/models/networks/flux2/infer/transformer_infer.py
@@ -32,6 +32,14 @@ def __init__(self, config):
             self.seq_p_fp4_comm = False
             self.enable_head_parallel = False
 
+    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 = block_weights.rope.apply(query, key, image_rotary_emb, positions=image_rotary_positions)
 
+        # 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,
             )
 
@@ -197,8 +211,13 @@ def infer_single_stream_block(
 
         query, key = block_weights.rope.apply(query, key, image_rotary_emb, positions=image_rotary_positions)
 
+        # 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")
 
@@ -208,7 +227,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,
@@ -222,10 +241,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 e94acbac6..440deab77 100644
--- a/lightx2v/models/networks/flux2/model.py
+++ b/lightx2v/models/networks/flux2/model.py
@@ -37,6 +37,10 @@ def __init__(self, config, model_path, device):
         self._init_tensor_parallel()
         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_tensor_parallel(self):
@@ -420,7 +424,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 db8d96aff..ab348ad40 100644
--- a/lightx2v/models/networks/flux2/weights/transformer_weights.py
+++ b/lightx2v/models/networks/flux2/weights/transformer_weights.py
@@ -258,8 +258,50 @@ def __init__(self, config):
         self.mm_type = config.get("dit_quant_scheme", "Default")
         self._configure_resident_blocks(config)
 
-        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 c0c92dc89..4abc3dc18 100644
--- a/lightx2v/models/runners/flux2/flux2_runner.py
+++ b/lightx2v/models/runners/flux2/flux2_runner.py
@@ -296,6 +296,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
 
@@ -320,6 +326,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],
@@ -428,13 +506,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 ddb0eb46e..271166dbd 100755
--- a/lightx2v/pipeline.py
+++ b/lightx2v/pipeline.py
@@ -429,11 +429,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 86b04fce0..a61be864f 100755
--- a/lightx2v/utils/set_config.py
+++ b/lightx2v/utils/set_config.py
@@ -27,6 +27,7 @@ def get_default_config():
         "parallel": False,
         "seq_parallel": False,
         "cfg_parallel": False,
+        "pipefusion_parallel": False,
         "enable_cfg": False,
         "warmup": False,
         "use_image_encoder": True,
@@ -413,11 +414,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})."
             )
 
         phase_aware = bool(config.get("model_cls") == "hunyuan_image3" and config["parallel"].get("phase_aware", False))
@@ -425,7 +427,10 @@ def set_parallel_config(config):
             from lightx2v.models.networks.hunyuan_image3.parallel import initialize_hunyuan_image3_parallel_runtime
 
             initialize_hunyuan_image3_parallel_runtime(config)
+            config["pipefusion_parallel"] = False
         elif tensor_p_size > 1:
+            if pp_size > 1:
+                raise ValueError("PipeFusion pipeline parallelism cannot be combined with tensor parallelism")
             # Tensor parallel is the innermost dimension. Optional CFG and
             # sequence dimensions are prepended so ranks with the same
             # non-TP coordinates form contiguous TP groups. For TP+SP+CFG:
@@ -448,13 +453,21 @@ def set_parallel_config(config):
             config["tensor_parallel"] = True
             config["seq_parallel"] = seq_p_size > 1
             config["cfg_parallel"] = bool(config.get("enable_cfg", False) and cfg_p_size > 1)
+            config["pipefusion_parallel"] = False
         else:
-            # Original 2D mesh for cfg_p and seq_p
-            config["device_mesh"] = init_device_mesh(AI_DEVICE, (cfg_p_size, seq_p_size), mesh_dim_names=("cfg_p", "seq_p"))
+            # Multi-dimensional mesh: (cfg_p, pp, seq_p)
+            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
             config["seq_parallel"] = seq_p_size > 1
             config["cfg_parallel"] = bool(config.get("enable_cfg", False) and cfg_p_size > 1)
 
+            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
         if AI_DEVICE == "cuda":
             warmup_device = f"{AI_DEVICE}:{torch.cuda.current_device()}"
@@ -462,6 +475,8 @@ def set_parallel_config(config):
             warmup_device = AI_DEVICE
         _a = torch.zeros([1], device=warmup_device)
         dist.all_reduce(_a)
+    else:
+        config["pipefusion_parallel"] = False
 
 
 def print_config(config):

From 199d5e82aa2a7c92827f4b04a0364f0f14e2283b 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/11] 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 43ad071030422674c9c8f4ed28e98ec4738b128d 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/11] 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 fe731a876ec816ea5438b2e990ec92c6404d524d 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/11] 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 e484e05f54345e7dadf0eaf211dbc085610aa401 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/11] 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 d515cbd7c73a0c9da3d1cc598a074eb7763c4afb 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/11] 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 2bfaa80dfd6e45525190184b9bb29067e6f924f0 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/11] 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 06eb1971aa11c53a5546b6c91b79c571d824b109 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/11] 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 4abc3dc18..44e4988da 100644
--- a/lightx2v/models/runners/flux2/flux2_runner.py
+++ b/lightx2v/models/runners/flux2/flux2_runner.py
@@ -341,12 +341,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 6681266bb84873f42d6e4f412f384520ad32a60d 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 09/11] 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 2d19874d6..98ecaead1 100755
--- a/lightx2v/common/ops/attn/flash_attn.py
+++ b/lightx2v/common/ops/attn/flash_attn.py
@@ -54,9 +54,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:
@@ -129,9 +130,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:
@@ -211,7 +213,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 440deab77..05857931f 100644
--- a/lightx2v/models/networks/flux2/model.py
+++ b/lightx2v/models/networks/flux2/model.py
@@ -284,7 +284,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 ab348ad40..ff778f20b 100644
--- a/lightx2v/models/networks/flux2/weights/transformer_weights.py
+++ b/lightx2v/models/networks/flux2/weights/transformer_weights.py
@@ -278,9 +278,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 44e4988da..6135c5414 100644
--- a/lightx2v/models/runners/flux2/flux2_runner.py
+++ b/lightx2v/models/runners/flux2/flux2_runner.py
@@ -356,6 +356,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 d1af6914c3f5bc59c7752f760f07046faa9118a0 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 10/11] 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 6135c5414..7c16837a8 100644
--- a/lightx2v/models/runners/flux2/flux2_runner.py
+++ b/lightx2v/models/runners/flux2/flux2_runner.py
@@ -509,6 +509,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 886187fa3b498b6bd76ba7d7b51e39b94a173150 Mon Sep 17 00:00:00 2001
From: Qin-sx <Qin_sx@outlook.com>
Date: Wed, 9 Sep 2026 23:33:43 +0800
Subject: [PATCH 11/11] update pipefusion

---
 lightx2v/common/distributed/__init__.py       |  1 +
 lightx2v/common/distributed/pipeline_comm.py  | 19 ++++--
 lightx2v/common/distributed/pipeline_state.py | 19 ++++++
 .../flux2/infer/pipefusion/pipeline_driver.py | 34 ++++++----
 .../infer/pipefusion/transformer_infer.py     | 37 ++++++-----
 .../networks/flux2/infer/transformer_infer.py |  9 +--
 .../flux2/weights/transformer_weights.py      |  2 +
 lightx2v/models/runners/flux2/flux2_runner.py | 19 +++++-
 lightx2v/models/schedulers/flux2/scheduler.py | 10 +++
 lightx2v/utils/set_config.py                  | 63 ++++++++++++++++---
 10 files changed, 165 insertions(+), 48 deletions(-)

diff --git a/lightx2v/common/distributed/__init__.py b/lightx2v/common/distributed/__init__.py
index 0f979d8f9..92b83ee0c 100644
--- a/lightx2v/common/distributed/__init__.py
+++ b/lightx2v/common/distributed/__init__.py
@@ -8,4 +8,5 @@
     init_pipeline_parallel_state,
     is_pipeline_first_stage,
     is_pipeline_last_stage,
+    reset_pipeline_parallel_state,
 )
diff --git a/lightx2v/common/distributed/pipeline_comm.py b/lightx2v/common/distributed/pipeline_comm.py
index 521913272..8b88174e7 100644
--- a/lightx2v/common/distributed/pipeline_comm.py
+++ b/lightx2v/common/distributed/pipeline_comm.py
@@ -15,7 +15,13 @@
 
 
 class PipelineComm:
-    """P2P communication between adjacent pipeline stages."""
+    """P2P communication between adjacent pipeline stages.
+
+    CUDA-only: receive buffers are allocated on the current CUDA device via
+    ``torch.cuda.current_device()``. PipeFusion is gated to CUDA at config time
+    (see ``set_config._validate_pipefusion_config``), so this is not reachable
+    on XPU/NPU platforms.
+    """
 
     def __init__(self, pp_group: dist.ProcessGroup):
         self.pp_group = pp_group
@@ -73,13 +79,14 @@ def pipeline_recv(self, name: str = "latent", shape=None, dtype=None) -> torch.T
     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.
+        Returns ``(Work, actual_send_tensor)``. ``actual_send_tensor`` is the
+        contiguous buffer actually handed to NCCL (``contiguous()`` may copy),
+        so the caller must keep BOTH alive until ``Work.wait()`` 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)
+        return dist.isend(tensor, dst=self.next_rank, group=self._device_group), tensor
 
     def add_pipeline_recv_task(self, idx: int = 0, name: str = "latent", shape=None, dtype=None):
         self._recv_tasks_queue.append((name, idx))
diff --git a/lightx2v/common/distributed/pipeline_state.py b/lightx2v/common/distributed/pipeline_state.py
index e84ba4b85..f7129c9d4 100644
--- a/lightx2v/common/distributed/pipeline_state.py
+++ b/lightx2v/common/distributed/pipeline_state.py
@@ -24,6 +24,18 @@ def init_pipeline_parallel_state(pp_group: dist.ProcessGroup):
     _runtime_state = PipelineRuntimeState()
 
 
+def reset_pipeline_parallel_state():
+    """Reset the pipeline-parallel global state.
+
+    Provides an explicit teardown so the same process can rebuild a runner,
+    switch parallel configs, or run multiple tests without carrying stale
+    group/state (see ``init_pipeline_parallel_state``).
+    """
+    global _pp_group, _runtime_state
+    _pp_group = None
+    _runtime_state = None
+
+
 # ---------------------------------------------------------------------------
 # Stage helpers
 # ---------------------------------------------------------------------------
@@ -109,6 +121,8 @@ def set_input_parameters(
         self.vae_scale_factor = vae_scale_factor
         self.patch_size = patch_size
         self.warmup_steps = warmup_steps
+        if warmup_steps <= 0:
+            raise ValueError(f"pipeline_warmup_steps must be >= 1, got {warmup_steps}.")
         if num_pipeline_patch is not None:
             self.num_pipeline_patch = num_pipeline_patch
 
@@ -123,6 +137,11 @@ def set_input_parameters(
             self.packed_w = width // multiple_of
             tok_count = self.packed_h * self.packed_w
 
+        if self.num_pipeline_patch <= 0:
+            raise ValueError(f"num_pipeline_patch must be >= 1, got {self.num_pipeline_patch}.")
+        if self.num_pipeline_patch > tok_count:
+            raise ValueError(f"num_pipeline_patch ({self.num_pipeline_patch}) exceeds token count ({tok_count}); each patch would be empty.")
+
         # Split tokens evenly across patches
         base = tok_count // self.num_pipeline_patch
         remainder = tok_count % self.num_pipeline_patch
diff --git a/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py b/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
index a512db1d3..948e0c92c 100644
--- a/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
+++ b/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py
@@ -19,6 +19,7 @@
     is_pipeline_first_stage,
     is_pipeline_last_stage,
 )
+from lightx2v.utils.envs import GET_DTYPE
 
 
 class Flux2PipelineDriver:
@@ -32,9 +33,7 @@ def __init__(self, model, config):
         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)
+        self._dtype = GET_DTYPE()
 
     # ==================================================================
     # Public entry point
@@ -45,6 +44,11 @@ def run_pipeline(self, latents, prompt_embeds, text_ids, latent_image_ids, times
 
         Returns final latents on the last stage, ``None`` on other stages.
         """
+        if do_cfg:
+            # PipeFusion does not maintain separate cond/uncond pipelines and
+            # KV caches, so CFG is not supported. The async path would silently
+            # drop CFG (see set_config validation, which rejects this earlier).
+            raise NotImplementedError("PipeFusion does not support CFG. Set sample_guide_scale <= 1.0 or enable_cfg=False.")
         warmup_steps = self.state.warmup_steps
 
         if self._pp_world_size > 1 and len(timesteps) > warmup_steps:
@@ -109,7 +113,7 @@ def _sync_pipeline(self, latents, prompt_embeds, text_ids, latent_image_ids, tim
                 uncond_result = self._sync_pass(
                     latents,
                     negative_prompt_embeds,
-                    negative_text_ids or text_ids,
+                    negative_text_ids if negative_text_ids is not None else text_ids,
                     latent_image_ids,
                     t,
                     scheduler,
@@ -301,6 +305,11 @@ def _async_pipeline(self, latents, prompt_embeds, text_ids, latent_image_ids, ti
                         patch_latents[patch_idx] = self.pp_comm.get_pipeline_recv_data(patch_idx, "latent")
 
                 # ---- 2. Compute (default stream) ----
+                # NOTE: encoder hidden states are only transferred once per
+                # timestep (patch 0) and reused for every subsequent patch.
+                # Flux double blocks update the text stream from the current
+                # image patch, so patch k>0 text states are a stale-text
+                # approximation — an intentional quality/memory trade-off.
                 cur_enc = prompt_embeds if self._is_first else last_encoder_hidden_states
                 result = self._async_backbone(
                     patch_latents[patch_idx],
@@ -311,21 +320,22 @@ def _async_pipeline(self, latents, prompt_embeds, text_ids, latent_image_ids, ti
                 )
 
                 # ---- 3. Send result (default stream, after compute) ----
-                # Store isend request to prevent tensor GC before send completes
+                # Store isend request AND the actual contiguous send buffer to
+                # prevent tensor GC before the send completes.
                 if self._is_last:
                     noise_pred = result
-                    scheduler.scheduler._step_index = i + self.state.warmup_steps
+                    scheduler.set_step_index(i + self.state.warmup_steps)
                     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]))
+                        req, sent = self.pp_comm.pipeline_isend(patch_latents[patch_idx], name="latent", segment_idx=patch_idx)
+                        pending_isends.append((req, sent))
                 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))
+                        req, sent = self.pp_comm.pipeline_isend(next_enc, name="encoder_hidden_state", segment_idx=0)
+                        pending_isends.append((req, sent))
+                    req, sent = self.pp_comm.pipeline_isend(hidden_states, name="latent", segment_idx=patch_idx)
+                    pending_isends.append((req, sent))
 
                 # ---- 4. Post next irecv (default stream — NCCL internal
                 # stream handles the actual async transfer; no cross-stream
diff --git a/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py b/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
index faad09e18..c0d76f5f3 100644
--- a/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
+++ b/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py
@@ -3,7 +3,10 @@
 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.
+   patches (stale-KV approximation), while text KV is recomputed from the
+   encoder hidden states it is given. NOTE: on non-first stages the async
+   driver feeds patch-0 text hidden states to every patch (a second,
+   stale-text approximation) — see ``pipeline_driver._async_pipeline``.
 3. Return ``(hidden_states, encoder_hidden_states, num_txt_tokens)`` so the
    pipeline driver can P2P-pass intermediate activations between stages.
 """
@@ -41,12 +44,13 @@ def __init__(self, config):
     # Stale-KV hook (overrides base class no-op)
     # ------------------------------------------------------------------
 
-    def _maybe_apply_stale_kv(self, key, value, num_txt_tokens, block_idx):
+    def _maybe_apply_stale_kv(self, key, value, num_txt_tokens, block_idx, block_type=None):
         """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.
+        - Cache is indexed by ((block_type, block_idx), patch_slot). Double and
+          single blocks both number their blocks from 0, so ``block_idx`` alone
+          would collide across the two types; ``block_type`` disambiguates.
         - 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
@@ -55,6 +59,7 @@ def _maybe_apply_stale_kv(self, key, value, num_txt_tokens, block_idx):
         Optimization: pre-allocated buffers + copy_ instead of torch.cat
         to avoid memory allocations per generation.
         """
+        cache_key = (block_type, block_idx)
         num_patch = self.pipeline_state.num_pipeline_patch
         if num_patch <= 1 or num_txt_tokens <= 0:
             return key, value
@@ -69,12 +74,12 @@ def _maybe_apply_stale_kv(self, key, value, num_txt_tokens, block_idx):
             # 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
+            if cache_key not in self._kv_cache:
+                self._kv_cache[cache_key] = [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] = [
+                self._kv_cache[cache_key][i] = [
                     split_ks[i].clone(),
                     split_vs[i].clone(),
                 ]
@@ -83,22 +88,22 @@ def _maybe_apply_stale_kv(self, key, value, num_txt_tokens, block_idx):
         # ---- 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
+        if cache_key not in self._kv_cache:
+            self._kv_cache[cache_key] = [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()]
+        self._kv_cache[cache_key][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 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)
+        if cache_key not in self._full_k_bufs or self._full_k_bufs[cache_key].shape[0] != full_len or self._full_k_bufs[cache_key].dtype != key.dtype:
+            self._full_k_bufs[cache_key] = torch.empty(full_len, *key.shape[1:], dtype=key.dtype, device=key.device)
+            self._full_v_bufs[cache_key] = torch.empty(full_len, *value.shape[1:], dtype=value.dtype, device=value.device)
 
-        buf_k = self._full_k_bufs[block_idx]
-        buf_v = self._full_v_bufs[block_idx]
+        buf_k = self._full_k_bufs[cache_key]
+        buf_v = self._full_v_bufs[cache_key]
 
         # Copy text K/V (fresh, from current patch's computation)
         buf_k[:num_txt_tokens].copy_(text_key)
@@ -114,7 +119,7 @@ def _maybe_apply_stale_kv(self, key, value, num_txt_tokens, block_idx):
                 buf_v[offset : offset + n].copy_(img_value)
             else:
                 # Stale from cache (previous timestep)
-                cached = self._kv_cache[block_idx][slot]
+                cached = self._kv_cache[cache_key][slot]
                 buf_k[offset : offset + n].copy_(cached[0])
                 buf_v[offset : offset + n].copy_(cached[1])
             offset += n
diff --git a/lightx2v/models/networks/flux2/infer/transformer_infer.py b/lightx2v/models/networks/flux2/infer/transformer_infer.py
index a37695a17..34e7ed84f 100644
--- a/lightx2v/models/networks/flux2/infer/transformer_infer.py
+++ b/lightx2v/models/networks/flux2/infer/transformer_infer.py
@@ -32,11 +32,12 @@ def __init__(self, config):
             self.seq_p_fp4_comm = False
             self.enable_head_parallel = False
 
-    def _maybe_apply_stale_kv(self, key, value, num_txt_tokens, block_idx):
+    def _maybe_apply_stale_kv(self, key, value, num_txt_tokens, block_idx, block_type=None):
         """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.
+        while keeping text KV fresh. ``block_type`` distinguishes double vs
+        single blocks, whose ``block_idx`` both restart from 0.
         """
         return key, value
 
@@ -105,7 +106,7 @@ def infer_double_stream_block(
 
         # 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)
+        key, value = self._maybe_apply_stale_kv(key, value, num_txt_tokens, block_weights.block_idx, block_type=block_weights.block_type)
 
         total_len = query.shape[0]
         kv_len = key.shape[0]  # may differ from total_len in PipeFusion (stale-KV)
@@ -212,7 +213,7 @@ def infer_single_stream_block(
         query, key = block_weights.rope.apply(query, key, image_rotary_emb, positions=image_rotary_positions)
 
         # 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)
+        key, value = self._maybe_apply_stale_kv(key, value, num_txt_tokens, block_weights.block_idx, block_type=block_weights.block_type)
 
         total_len = query.shape[0]
         kv_len = key.shape[0]  # may differ from total_len in PipeFusion (stale-KV)
diff --git a/lightx2v/models/networks/flux2/weights/transformer_weights.py b/lightx2v/models/networks/flux2/weights/transformer_weights.py
index ff778f20b..300646a25 100644
--- a/lightx2v/models/networks/flux2/weights/transformer_weights.py
+++ b/lightx2v/models/networks/flux2/weights/transformer_weights.py
@@ -139,6 +139,7 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe
         super().__init__()
         self.config = config
         self.block_idx = block_idx
+        self.block_type = "double"
         self.inner_dim = config["num_attention_heads"] * config["attention_head_dim"]
         self.mm_type = config.get("dit_quant_scheme", "Default")
         self.layer_norm_type = config.get("layer_norm_type", "torch")
@@ -206,6 +207,7 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe
         super().__init__()
         self.config = config
         self.block_idx = block_idx
+        self.block_type = "single"
         self.inner_dim = config["num_attention_heads"] * config["attention_head_dim"]
         self.mm_type = config.get("dit_quant_scheme", "Default")
         self.layer_norm_type = config.get("layer_norm_type", "torch")
diff --git a/lightx2v/models/runners/flux2/flux2_runner.py b/lightx2v/models/runners/flux2/flux2_runner.py
index 7c16837a8..e4d0905ff 100644
--- a/lightx2v/models/runners/flux2/flux2_runner.py
+++ b/lightx2v/models/runners/flux2/flux2_runner.py
@@ -508,14 +508,27 @@ def run_pipeline(self, input_info):
         if self.config.get("pipefusion_parallel", False):
             from lightx2v.common.distributed import is_pipeline_last_stage
 
+            if input_info.return_result_tensor:
+                # Final latents/images exist only on the last pipeline stage and
+                # there is no cross-rank gather implemented, so rank 0 cannot
+                # return them under the standard tensor-return contract.
+                raise NotImplementedError("PipeFusion does not support return_result_tensor yet; the result exists only on the last pipeline stage.")
+
+            # Clear the stale-KV cache on EVERY rank between requests. Only the
+            # last stage runs VAE decode, but each stage holds its own cache and
+            # must not carry stale KV / full K-V buffers into the next request.
+            if hasattr(self.model.transformer_infer, "clear_kv_cache"):
+                self.model.transformer_infer.clear_kv_cache()
+
             if is_pipeline_last_stage():
-                # Offload transformer weights and clear KV cache before VAE decode to avoid OOM
+                # Offload transformer weights before VAE decode to avoid OOM,
+                # then move them back afterwards so a resident runner (serving)
+                # can process the next request with weights on the device.
                 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)
+                self.model.transformer_weights.to_cuda()
             else:
                 images = None
         else:
diff --git a/lightx2v/models/schedulers/flux2/scheduler.py b/lightx2v/models/schedulers/flux2/scheduler.py
index 432bf4321..5183e251c 100755
--- a/lightx2v/models/schedulers/flux2/scheduler.py
+++ b/lightx2v/models/schedulers/flux2/scheduler.py
@@ -146,6 +146,16 @@ def step_post(self):
         )
         self.latents = latents
 
+    def set_step_index(self, step_index):
+        """Set the diffusers scheduler's internal step index.
+
+        The async PipeFusion driver steps patch-by-patch without going through
+        ``step_pre``/``step_post`` for every patch, so the underlying
+        scheduler's ``_step_index`` (used by ``step()`` to look up the correct
+        sigma) must be advanced explicitly.
+        """
+        self.scheduler._step_index = step_index
+
     def step_post_patch(self, noise_pred, latents, t):
         """Patch-level scheduler step for async PipeFusion mode.
 
diff --git a/lightx2v/utils/set_config.py b/lightx2v/utils/set_config.py
index a61be864f..c9560952d 100755
--- a/lightx2v/utils/set_config.py
+++ b/lightx2v/utils/set_config.py
@@ -409,6 +409,45 @@ def set_config(args):
     return config
 
 
+def _validate_pipefusion_config(config):
+    """Reject unsupported PipeFusion combinations instead of silently misbehaving.
+
+    PipeFusion is currently a narrow feature: Flux2 Klein, T2I only, CUDA only,
+    no CFG, and no stacking with SP / TP / feature-caching / cpu-offload. The
+    pipeline driver only implements that slice; anything else must fail loudly
+    at config time rather than run incorrectly (e.g. dropping CFG) or crash.
+    """
+    model_cls = config.get("model_cls")
+    is_klein = model_cls == "flux2_klein" or (model_cls == "flux2" and config.get("model_variant") == "klein")
+    if not is_klein:
+        raise ValueError(
+            "PipeFusion is only supported for the Flux2 Klein model "
+            "(model_cls='flux2_klein', or model_cls='flux2' with model_variant='klein'); "
+            f"got model_cls={model_cls!r}, model_variant={config.get('model_variant')!r}."
+        )
+    if config.get("task", "t2i") != "t2i":
+        raise ValueError(f"PipeFusion currently supports only the 't2i' task, got {config.get('task', 't2i')!r}.")
+    if AI_DEVICE != "cuda":
+        raise ValueError(f"PipeFusion requires CUDA, but AI_DEVICE={AI_DEVICE!r}.")
+    if config.get("feature_caching", "NoCaching") not in ("NoCaching", "None"):
+        raise ValueError(f"PipeFusion cannot be combined with feature_caching={config.get('feature_caching')!r}.")
+    if config.get("cpu_offload", False):
+        raise ValueError("PipeFusion cannot be combined with cpu_offload.")
+    if config.get("enable_cfg", False) and config.get("sample_guide_scale", 1.0) > 1.0:
+        raise ValueError("PipeFusion does not support CFG; set sample_guide_scale <= 1.0 or enable_cfg=False.")
+    if config["parallel"].get("seq_p_size", 1) > 1:
+        raise ValueError("PipeFusion cannot be combined with sequence parallel (seq_p_size > 1).")
+    if config["parallel"].get("cfg_p_size", 1) > 1:
+        raise ValueError("PipeFusion cannot be combined with CFG parallel (cfg_p_size > 1).")
+
+    num_patch = int(config["parallel"].get("num_pipeline_patch", 4))
+    if num_patch <= 0:
+        raise ValueError(f"num_pipeline_patch must be >= 1, got {num_patch}.")
+    warmup_steps = int(config["parallel"].get("pipeline_warmup_steps", 1))
+    if warmup_steps <= 0:
+        raise ValueError(f"pipeline_warmup_steps must be >= 1, got {warmup_steps}.")
+
+
 def set_parallel_config(config):
     if config["parallel"]:
         tensor_p_size = int(config["parallel"].get("tensor_p_size", 1))
@@ -455,18 +494,28 @@ def set_parallel_config(config):
             config["cfg_parallel"] = bool(config.get("enable_cfg", False) and cfg_p_size > 1)
             config["pipefusion_parallel"] = False
         else:
-            # Multi-dimensional mesh: (cfg_p, pp, seq_p)
-            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
-            config["seq_parallel"] = seq_p_size > 1
-            config["cfg_parallel"] = bool(config.get("enable_cfg", False) and cfg_p_size > 1)
-
-            config["pipefusion_parallel"] = pp_size > 1
             if pp_size > 1:
+                # PipeFusion pipeline parallelism: 3D mesh (cfg_p, pp, seq_p).
+                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
+                config["seq_parallel"] = seq_p_size > 1
+                config["cfg_parallel"] = bool(config.get("enable_cfg", False) and cfg_p_size > 1)
+
+                config["pipefusion_parallel"] = True
+                _validate_pipefusion_config(config)
                 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)
+            else:
+                # No pipeline parallelism: keep the legacy 2D mesh (cfg_p, seq_p)
+                # unchanged so existing SP/CFG paths don't change shape or rank
+                # layout when pp_size == 1.
+                config["device_mesh"] = init_device_mesh(AI_DEVICE, (cfg_p_size, seq_p_size), mesh_dim_names=("cfg_p", "seq_p"))
+                config["tensor_parallel"] = False
+                config["seq_parallel"] = seq_p_size > 1
+                config["cfg_parallel"] = bool(config.get("enable_cfg", False) and cfg_p_size > 1)
+                config["pipefusion_parallel"] = False
 
         # warmup dist
         if AI_DEVICE == "cuda":