0 Fork 0

Files

T

Rustam Khadipash dc6177520e [Feature] add Tensor Parallelism to LongCat-Image(-Edit) (#926 )

Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>

2026-02-02 10:28:25 +00:00

22 KiB

Raw Blame History

Parallelism Acceleration Guide

This guide includes how to use parallelism methods in vLLM-Omni to speed up diffusion model inference as well as reduce the memory requirement on each device.

Overview

The following parallelism methods are currently supported in vLLM-Omni:

DeepSpeed Ulysses Sequence Parallel (DeepSpeed Ulysses-SP) (arxiv paper): Ulysses-SP splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.
Ring-Attention - splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results, keeping the sequence dimension sharded
Classifier-Free-Guidance Parallel (CFG-Parallel): CFG-Parallel runs the positive/negative prompts of classifier-free guidance (CFG) on different devices, then merges on a single device to perform the scheduler step.
Tensor Parallelism: Tensor parallelism shards model weights across devices. This can reduce per-GPU memory usage. Note that for diffusion models we currently shard the majority of layers within the DiT.

The following table shows which models are currently supported by parallelism method:

ImageGen

Model	Model Identifier	Ulysses-SP	Ring-SP	CFG-Parallel	Tensor-Parallel
LongCat-Image	`meituan-longcat/LongCat-Image`	✅	✅	❌	✅
LongCat-Image-Edit	`meituan-longcat/LongCat-Image-Edit`	✅	✅	❌	✅
Ovis-Image	`OvisAI/Ovis-Image`	❌	❌	❌	❌
Qwen-Image	`Qwen/Qwen-Image`	✅	✅	✅	✅
Qwen-Image-Edit	`Qwen/Qwen-Image-Edit`	✅	✅	✅	✅
Qwen-Image-Edit-2509	`Qwen/Qwen-Image-Edit-2509`	✅	✅	✅	✅
Qwen-Image-Layered	`Qwen/Qwen-Image-Layered`	✅	✅	✅	✅
Z-Image	`Tongyi-MAI/Z-Image-Turbo`	✅	✅	❌	✅ (TP=2 only)
Stable-Diffusion3.5	`stabilityai/stable-diffusion-3.5`	❌	❌	❌	❌
FLUX.2-klein	`black-forest-labs/FLUX.2-klein-4B`	❌	❌	❌	✅
FLUX.1-dev	`black-forest-labs/FLUX.1-dev`	❌	❌	❌	✅

!!! note "TP Limitations for Diffusion Models" We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the text_encoder component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.

- Good news: The text_encoder typically has minimal impact on overall inference performance.
- Bad news: When TP is enabled, every TP process retains a full copy of the text_encoder weights, leading to significant GPU memory waste.

We are actively refactoring this design to address this. For details and progress, please refer to [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).

!!! note "Why Z-Image is TP=2 only" Z-Image Turbo is currently limited to tensor_parallel_size of 1 or 2 due to model shape divisibility constraints. For example, the model has n_heads=30 and a final projection out dimension of 64, so valid TP sizes must divide both 30 and 64; the only common divisors are 1 and 2.

VideoGen

Model	Model Identifier	Ulysses-SP	Ring-SP	Tensor-Parallel
Wan2.2	`Wan-AI/Wan2.2-T2V-A14B-Diffusers`	✅	✅	❌

Tensor Parallelism

Tensor parallelism splits model parameters across GPUs. In vLLM-Omni, tensor parallelism is configured via DiffusionParallelConfig.tensor_parallel_size.

Offline Inference

from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig

omni = Omni(
    model="Tongyi-MAI/Z-Image-Turbo",
    parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
)

outputs = omni.generate(
    "a cat reading a book",
    OmniDiffusionSamplingParams(
        num_inference_steps=9,
        width=512,
        height=512,
    ),
)

Sequence Parallelism

Ulysses-SP

Offline Inference

An example of offline inference script using Ulysses-SP is shown below:

from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
ulysses_degree = 2

omni = Omni(
    model="Qwen/Qwen-Image",
    parallel_config=DiffusionParallelConfig(ulysses_degree=2)
)

outputs = omni.generate(
    "A cat sitting on a windowsill",
    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
)

See examples/offline_inference/text_to_image/text_to_image.py for a complete working example.

Online Serving

You can enable Ulysses-SP in online serving for diffusion models via --usp:

# Text-to-image (requires >= 2 GPUs)
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2

Benchmarks

!!! note "Benchmark Disclaimer" These benchmarks are provided for general reference only. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:

- Specific model and use case
- Hardware configuration
- Careful parameter tuning
- Different inference settings (e.g., number of steps, image resolution)

To measure the parallelism methods, we run benchmarks with Qwen/Qwen-Image model generating images (2048x2048 as long sequence input) with 50 inference steps. The hardware devices are NVIDIA H800 GPUs. sdpa is the attention backends.

Configuration	Ulysses degree	Generation Time	Speedup
Baseline (diffusers)	-	112.5s	1.0x
Ulysses-SP	2	65.2s	1.73x
Ulysses-SP	4	39.6s	2.84x
Ulysses-SP	8	30.8s	3.65x

Ring-Attention

Ring-Attention (arxiv paper) splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results. Unlike Ulysses-SP which uses all-to-all communication, Ring-Attention keeps the sequence dimension sharded throughout the computation and circulates Key/Value blocks through a ring topology.

Offline Inference

An example of offline inference script using Ring-Attention is shown below:

from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
ring_degree = 2

omni = Omni(
    model="Qwen/Qwen-Image",
    parallel_config=DiffusionParallelConfig(ring_degree=2)
)

outputs = omni.generate(
    "A cat sitting on a windowsill",
    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
)

See examples/offline_inference/text_to_image/text_to_image.py for a complete working example.

Online Serving

You can enable Ring-Attention in online serving for diffusion models via --ring:

# Text-to-image (requires >= 2 GPUs)
vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2

Benchmarks

- Specific model and use case
- Hardware configuration
- Careful parameter tuning
- Different inference settings (e.g., number of steps, image resolution)

To measure the parallelism methods, we run benchmarks with Qwen/Qwen-Image model generating images (1024x1024 as long sequence input) with 50 inference steps. The hardware devices are NVIDIA A100 GPUs. flash_attn is the attention backends.

Configuration	Ring degree	Generation Time	Speedup
Baseline (diffusers)	-	45.2s	1.0x
Ring-Attention	2	29.9s	1.51x
Ring-Attention	4	23.3s	1.94x

Hybrid Ulysses + Ring

You can combine both Ulysses-SP and Ring-Attention for larger scale parallelism. The total sequence parallel size equals ulysses_degree × ring_degree.

Offline Inference

from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig

# Hybrid: 2 Ulysses × 2 Ring = 4 GPUs total
omni = Omni(
    model="Qwen/Qwen-Image",
    parallel_config=DiffusionParallelConfig(ulysses_degree=2, ring_degree=2)
)

outputs = omni.generate(
    "A cat sitting on a windowsill",
    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
)

Online Serving

# Text-to-image (requires >= 4 GPUs)
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2

Benchmarks

- Specific model and use case
- Hardware configuration
- Careful parameter tuning
- Different inference settings (e.g., number of steps, image resolution)

Configuration	Ulysses degree	Ring degree	Generation Time	Speedup
Baseline (diffusers)	-	-	45.2s	1.0x
Hybrid Ulysses + Ring	2	2	24.3s	1.87x

How to parallelize a new model

NOTE: "Terminology: SP vs CP" Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in the diffusers library. We use "Sequence Parallelism" to align with vLLM-Omni's terminology.

Non-intrusive `_sp_plan` (Recommended)

The _sp_plan mechanism allows SP without modifying forward() logic. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries.

Requirements for forward() function:

Tensor operations that need sharding/gathering must happen at nn.Module boundaries (not inline Python operations)
If your forward() contains inline tensor operations (e.g., torch.cat, pad_sequence) that need sharding, extract them into a submodule

When to create a submodule:

# ❌ BAD: Inline operations - hooks cannot intercept
def forward(self, x, cap_feats):
    unified = torch.cat([x, cap_feats], dim=1)  # Cannot be sharded via _sp_plan
    ...

# ✅ GOOD: Extract into a submodule
class UnifiedPrepare(nn.Module):
    def forward(self, x, cap_feats):
        return torch.cat([x, cap_feats], dim=1)  # Now can be sharded via _sp_plan

class MyModel(nn.Module):
    def __init__(self):
        self.unified_prepare = UnifiedPrepare()  # Submodule

    def forward(self, x, cap_feats):
        unified = self.unified_prepare(x, cap_feats)  # Hook can intercept here

Defining `_sp_plan`

Type definitions (see diffusers _modeling_parallel.py for reference):

from vllm_omni.diffusion.distributed.sp_plan import (
    SequenceParallelInput,   # Corresponds to diffusers' ContextParallelInput
    SequenceParallelOutput,  # Corresponds to diffusers' ContextParallelOutput
)

Parameter	Description
`split_dim`	Dimension to split/gather (usually `1` for sequence)
`expected_dims`	Expected tensor rank for validation (optional)
`split_output`	`False`: shard input parameters; `True`: shard output tensors
`auto_pad`	Auto-pad if sequence not divisible by world_size (Ulysses only)

Key naming convention:

Key	Meaning	Python equivalent
`""`	Root model	`model`
`"blocks.0"`	First element of ModuleList	`model.blocks[0]`
`"blocks.*"`	All elements of ModuleList	`for b in model.blocks`
`"outputs.main"`	ModuleDict entry	`model.outputs["main"]`

Dictionary key types:

Key type	`split_output`	Description
`"param_name"` (str)	`False`	Shard input parameter by name
`0`, `1` (int)	`True`	Shard output tuple by index

Example (similar to diffusers transformer_wan.py):

class MyTransformer(nn.Module):
    _sp_plan = {
        # Shard rope module OUTPUTS (returns tuple)
        "rope": {
            0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # cos
            1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # sin
        },
        # Shard transformer block INPUT parameter
        "blocks.0": {
            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
        },
        # Gather at final projection
        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
    }

Hook flow

Input → [SequenceParallelSplitHook: pre_forward] → Module.forward() → [post_forward] → ...
                                                                              ↓
... → [SequenceParallelGatherHook: post_forward] → Output

SplitHook shards tensors before/after the target module
Attention layers handle Ulysses/Ring communication internally
GatherHook collects sharded outputs

The framework automatically applies these hooks when sequence_parallel_size > 1.

Method 2: Intrusive modification (For complex cases)

For models with dynamic sharding logic that cannot be expressed via _sp_plan:

from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather

def forward(self, hidden_states, ...):
    if self.parallel_config.sequence_parallel_size > 1:
        hidden_states = sp_shard(hidden_states, dim=1)
        # ... computation ...
        output = sp_gather(output, dim=1)
    return output

Choosing the right approach

Scenario	Approach
Standard transformer	`_sp_plan`
Inline tensor ops need sharding	Extract to submodule + `_sp_plan`
Dynamic/conditional sharding	Intrusive modification

CFG-Parallel

Offline Inference

CFG-Parallel is enabled through DiffusionParallelConfig(cfg_parallel_size=2), which runs one rank for the positive branch and one rank for the negative branch.

An example of offline inference using CFG-Parallel (image-to-image) is shown below:

from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig

image_path = "path_to_image.png"
omni = Omni(
    model="Qwen/Qwen-Image-Edit",
    parallel_config=DiffusionParallelConfig(cfg_parallel_size=2),
)
input_image = Image.open(image_path).convert("RGB")

outputs = omni.generate(
    {
        "prompt": "turn this cat to a dog",
        "negative_prompt": "low quality, blurry",
        "multi_modal_data": {"image": input_image},
    },
    OmniDiffusionSamplingParams(
        true_cfg_scale=4.0,
        num_inference_steps=50,
    ),
)

Notes:

CFG-Parallel is only effective when a negative_prompt is provided AND a guidance scale (or cfg_scale) is greater than 1.

See examples/offline_inference/image_to_image/image_edit.py for a complete working example.

cd examples/offline_inference/image_to_image/
python image_edit.py \
  --model "Qwen/Qwen-Image-Edit" \
  --image "qwen_image_output.png" \
  --prompt "turn this cat to a dog" \
  --negative_prompt "low quality, blurry" \
  --cfg_scale 4.0 \
  --output "edited_image.png" \
  --cfg_parallel_size 2

Online Serving

You can enable CFG-Parallel in online serving for diffusion models via --cfg-parallel-size:

vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2

How to parallelize a pipeline

This section describes how to add CFG-Parallel to a diffusion pipeline. We use the Qwen-Image pipeline (vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py) as the reference implementation.

In QwenImagePipeline, each diffusion step runs two denoiser forward passes sequentially:

positive (prompt-conditioned)
negative (negative-prompt-conditioned)

CFG-Parallel assigns these two branches to different ranks in the CFG group and synchronizes the results.

vLLM-omni provides CFGParallelMixin base class that encapsulates the CFG parallel logic. By inheriting from this mixin and calling its methods, pipelines can easily implement CFG parallel without writing repetitive code.

Key Methods in CFGParallelMixin:

predict_noise_maybe_with_cfg(): Automatically handles CFG parallel noise prediction
scheduler_step_maybe_with_cfg(): Scheduler step with automatic CFG rank synchronization

Example Implementation:

class QwenImageCFGParallelMixin(CFGParallelMixin):
    """
    Base Mixin class for Qwen Image pipelines providing shared CFG methods.
    """

    def diffuse(
        self,
        prompt_embeds: torch.Tensor,
        prompt_embeds_mask: torch.Tensor,
        negative_prompt_embeds: torch.Tensor,
        negative_prompt_embeds_mask: torch.Tensor,
        latents: torch.Tensor,
        img_shapes: torch.Tensor,
        txt_seq_lens: torch.Tensor,
        negative_txt_seq_lens: torch.Tensor,
        timesteps: torch.Tensor,
        do_true_cfg: bool,
        guidance: torch.Tensor,
        true_cfg_scale: float,
        image_latents: torch.Tensor | None = None,
        cfg_normalize: bool = True,
        additional_transformer_kwargs: dict[str, Any] | None = None,
    ) -> torch.Tensor:
        self.transformer.do_true_cfg = do_true_cfg

        for i, t in enumerate(timesteps):
            timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)

            # Prepare kwargs for positive (conditional) prediction
            positive_kwargs = {
                "hidden_states": latents,
                "timestep": timestep / 1000,
                "guidance": guidance,
                "encoder_hidden_states_mask": prompt_embeds_mask,
                "encoder_hidden_states": prompt_embeds,
                "img_shapes": img_shapes,
                "txt_seq_lens": txt_seq_lens,
            }

            # Prepare kwargs for negative (unconditional) prediction
            if do_true_cfg:
                negative_kwargs = {
                    "hidden_states": latents,
                    "timestep": timestep / 1000,
                    "guidance": guidance,
                    "encoder_hidden_states_mask": negative_prompt_embeds_mask,
                    "encoder_hidden_states": negative_prompt_embeds,
                    "img_shapes": img_shapes,
                    "txt_seq_lens": negative_txt_seq_lens,
                }
            else:
                negative_kwargs = None

            # Predict noise with automatic CFG parallel handling
            # - In CFG parallel mode: rank0 computes positive, rank1 computes negative
            # - Automatically gathers results and combines them on rank0
            noise_pred = self.predict_noise_maybe_with_cfg(
                do_true_cfg=do_true_cfg,
                true_cfg_scale=true_cfg_scale,
                positive_kwargs=positive_kwargs,
                negative_kwargs=negative_kwargs,
                cfg_normalize=cfg_normalize,
            )

            # Step scheduler with automatic CFG synchronization
            # - Only rank0 computes the scheduler step
            # - Automatically broadcasts updated latents to all ranks
            latents = self.scheduler_step_maybe_with_cfg(
                noise_pred, t, latents, do_true_cfg
            )

        return latents

How it works:

Prepare separate positive_kwargs and negative_kwargs for conditional and unconditional predictions
Call predict_noise_maybe_with_cfg() which:
- Detects if CFG parallel is enabled (get_classifier_free_guidance_world_size() > 1)
- Distributes computation: rank0 processes positive, rank1 processes negative
- Gathers predictions and combines them using combine_cfg_noise() on rank0
- Returns combined noise prediction (only valid on rank0)
Call scheduler_step_maybe_with_cfg() which:
- Only rank0 computes the scheduler step
- Broadcasts the updated latents to all ranks for synchronization

How to customize

Some pipelines may need to customize the following functions in CFGParallelMixin:

You may need to edit predict_noise function for custom behaviors.

def predict_noise(self, *args, **kwargs):
    """
    Forward pass through transformer to predict noise.

    Subclasses should override this if they need custom behavior,
    but the default implementation calls self.transformer.
    """
    return self.transformer(*args, **kwargs)[0]

The default normalization function after combining the noise predictions from both branches is as follows. You may need to customize it.

def cfg_normalize_function(self, noise_pred, comb_pred):
    """
    Normalize the combined noise prediction.

    Args:
        noise_pred: positive noise prediction
        comb_pred: combined noise prediction after CFG

    Returns:
        Normalized noise prediction tensor
    """
    cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
    noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
    noise_pred = comb_pred * (cond_norm / noise_norm)
    return noise_pred

22 KiB Raw Blame History Unescape Escape

Parallelism Acceleration Guide

Overview

ImageGen

VideoGen

Tensor Parallelism

Offline Inference

Sequence Parallelism

Ulysses-SP

Offline Inference

Online Serving

Benchmarks

Ring-Attention

Offline Inference

Online Serving

Benchmarks

Hybrid Ulysses + Ring

Offline Inference

Online Serving

Benchmarks

How to parallelize a new model

Non-intrusive _sp_plan (Recommended)

Defining _sp_plan

Hook flow

Method 2: Intrusive modification (For complex cases)

Choosing the right approach

CFG-Parallel

Offline Inference

Online Serving

How to parallelize a pipeline

22 KiB

Raw Blame History

Non-intrusive `_sp_plan` (Recommended)

Defining `_sp_plan`