feat: support pipefusion for flux2

[Qin-sx]

Support PipeFusion parallel inference for Flux2

Summary

Adds pipeline parallelism (PipeFusion) support for the Flux2 model in LightX2V, enabling multi-GPU inference with per-patch stale-KV caching across pipeline stages.

Key Components

New modules

• lightx2v/common/distributed/ — pipeline parallel infrastructure:
  • pipeline_comm.py (PipelineComm): P2P communication manager between adjacent stages. Wraps dist.isend/dist.irecv with pre-allocated recv buffers, a FIFO task queue, and async recv_next() prefetch.
  • pipeline_state.py (PipelineRuntimeState): patch metadata management and stage-identification helpers (is_pipeline_first_stage, get_pipeline_parallel_rank, etc.).
• lightx2v/models/networks/flux2/infer/pipefusion/ — Flux2 PipeFusion driver:
  • pipeline_driver.py (Flux2PipelineDriver): orchestrates the denoising loop across stages — sync pipeline for warmup, async pipeline for the main loop with per-patch compute/comm overlap.
  • transformer_infer.py (Flux2PipeFusionTransformerInfer): subclasses Flux2TransformerInfer to run only the current stage's block subset and apply stale-KV caching (image KV cached across patches, text KV stays fresh).

Comparison

Stage	PP8	USP8
Load models	5.38 ~ 7.68s	8.95 ~ 9.01s
Run Text Encoder	0.43 ~ 0.52s	0.45 ~ 0.49s
Run Encoders (incl. above)	0.68 ~ 0.84s	0.72 ~ 0.79s
DiT denoising loop	8.58 ~ 10.93s	7.96 ~ 8.08s
RUN pipeline (excl. load, incl. VAE)	9.84 ~ 12.06s (normal ranks) 14.40s (rank7, incl. VAE)	10.56 ~ 10.75s (normal ranks) 11.62s (rank0, incl. VAE)
Run VAE Decoder	1.04s	1.05 ~ 1.11s
Total Cost (incl. model load)	17.83 ~ 18.05s (normal ranks) 20.05s (rank7, slowest)	19.86 ~ 20.02s (normal ranks) 20.92s (rank0, slowest)

PP

{
    "model_cls": "flux2_klein",
    "task": "t2i",
    "infer_steps": 20,
    "sample_guide_scale": 1.0,
    "vae_scale_factor": 16,
    "feature_caching": "None",
    "enable_cfg": false,
    "patch_size": 2,
    "tokenizer_max_length": 512,
    "rope_type": "flashinfer_rope",
    "attn_type": "flash_attn2",
    "parallel": {
        "cfg_p_size": 1,
        "pp_size": 8,
        "seq_p_size": 1,
        "num_pipeline_patch": 8,
        "pipeline_warmup_steps": 2
    }
}

torchrun --nproc_per_node=8 -m lightx2v.infer \
    --model_cls flux2_klein \
    --task t2i \
    --target_shape 2048 2048 \
    --model_path $model_path \
    --prompt "A cat running in a garden" \
    --save_result_path "save_results/flux2_klein_4b_pp8.png" \
    --config_json "${lightx2v_path}/configs/flux2/flux2_klein_pp8.json" \
    2>&1 | tee "$log_file"

USP

{
    "model_cls": "flux2_klein",
    "task": "t2i",
    "infer_steps": 20,
    "sample_guide_scale": 1.0,
    "vae_scale_factor": 16,
    "feature_caching": "None",
    "enable_cfg": false,
    "patch_size": 2,
    "tokenizer_max_length": 512,
    "rope_type": "flashinfer_rope",
    "attn_type": "flash_attn2",
    "parallel": {
        "seq_p_size": 8,
        "seq_p_attn_type": "ulysses"
    }
}

torchrun --nproc_per_node=8 -m lightx2v.infer \
    --model_cls flux2_klein \
    --task t2i \
    --target_shape 2048 2048 \
    --model_path $model_path \
    --prompt "A cat running in a garden" \
    --save_result_path "save_results/flux2_klein_4b_usp8.png" \
    --config_json "${lightx2v_path}/configs/flux2/flux2_klein_usp8.json" \
    2>&1 | tee "$log_file"

[gemini-code-assist]

Code Review

This pull request introduces PipeFusion-based pipeline parallelism for Flux2, adding a pipeline driver, stale-KV caching transformer inference, and distributed communication helpers to orchestrate sync and async denoising loops. Key feedback from the review highlights critical data corruption risks due to globally shared pre-allocated K/V buffers across transformer blocks, which should be stored in block-specific dictionaries instead. Additionally, the reviewer suggested performance and code quality improvements, including increasing the pending isend queue limit to prevent CPU blocking, avoiding redundant memory allocations and CPU-GPU synchronization bottlenecks, and utilizing the newly added step_post_patch scheduler method for better encapsulation.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.