vllm-omni/.claude/skills/add-diffusion-model/references/parallelism-patterns.md

Parallelism Patterns Reference

Overview

vLLM-Omni supports multiple parallelism strategies for diffusion models. Each targets a different bottleneck:

Strategy	Splits	Best For	Constraint
Tensor Parallel (TP)	Model layers across GPUs	Latency reduction, large models	Requires fast GPU interconnect, num_heads % tp == 0
Sequence Parallel (SP/USP)	Sequence tokens across GPUs	Long sequences (video, high-res)	Near-linear scaling
CFG Parallel	Positive/negative CFG branches	Models using classifier-free guidance	Exactly 2 GPUs
HSDP	Weight shards via FSDP2	VRAM reduction	Cannot combine with TP
VAE Patch Parallel	VAE decode spatial tiles	Large VAE outputs	Auto-enables tiling

Recommended integration order: TP → SP → CFG Parallel → HSDP

Official design docs:

• TP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/tensor_parallel
• SP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/sequence_parallel
• CFG: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cfg_parallel
• HSDP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/hsdp

---

Tensor Parallelism (TP)

Replace standard nn.Linear with vLLM's parallel linear layers. This is the most invasive change but provides direct VRAM savings and compute speedup.

Layer replacement rules

Pattern	vLLM Layer	When to Use
Fan-out (first in FFN)	ColumnParallelLinear	Projection that splits output across ranks
Fan-in (second in FFN)	RowParallelLinear	Projection that gathers across ranks
QKV projection	QKVParallelLinear	Fused Q/K/V for self-attention
Single Q or K or V	ColumnParallelLinear	Separate projections (cross-attention)
Attention output	RowParallelLinear	Output projection after attention
Must not shard	ReplicatedLinear	Layers that must stay replicated

MLP Block (Up-Down Pattern)

from vllm.model_executor.layers.linear import (
    ColumnParallelLinear, RowParallelLinear,
)

class TPFeedForward(nn.Module):
    def __init__(self, dim, ffn_dim):
        super().__init__()
        self.fc1 = ColumnParallelLinear(dim, ffn_dim, bias=False, return_bias=False)
        self.fc2 = RowParallelLinear(
            ffn_dim, dim, bias=False,
            input_is_parallel=True,  # Input already sharded from fc1
            return_bias=False,
        )

    def forward(self, x):
        x, _ = self.fc1(x)
        x = torch.nn.functional.gelu(x)
        x, _ = self.fc2(x)
        return x

Attention Block (QKV-Out Pattern)

from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
from vllm_omni.diffusion.attention.layer import Attention

class TPSelfAttention(nn.Module):
    def __init__(self, dim, num_heads, num_kv_heads=None):
        super().__init__()
        num_kv_heads = num_kv_heads or num_heads
        self.head_dim = dim // num_heads

        self.to_qkv = QKVParallelLinear(
            hidden_size=dim,
            head_size=self.head_dim,
            total_num_heads=num_heads,
            total_num_kv_heads=num_kv_heads,
            bias=False,
            return_bias=False,
        )
        self.to_out = RowParallelLinear(
            dim, dim, bias=False,
            input_is_parallel=True,
            return_bias=False,
        )
        self.attn = Attention(
            num_heads=self.to_qkv.num_heads,      # Local heads per GPU
            head_size=self.head_dim,
            softmax_scale=1.0 / (self.head_dim ** 0.5),
            causal=False,
            num_kv_heads=self.to_qkv.num_kv_heads,  # Local KV heads per GPU
        )

    def forward(self, x):
        qkv, _ = self.to_qkv(x)
        q, k, v = qkv.split(
            [self.to_qkv.num_heads * self.head_dim,
             self.to_qkv.num_kv_heads * self.head_dim,
             self.to_qkv.num_kv_heads * self.head_dim],
            dim=-1,
        )
        B, S, _ = x.shape
        q = q.view(B, S, self.to_qkv.num_heads, self.head_dim)
        k = k.view(B, S, self.to_qkv.num_kv_heads, self.head_dim)
        v = v.view(B, S, self.to_qkv.num_kv_heads, self.head_dim)
        out = self.attn(q, k, v)
        out = out.reshape(B, S, -1)
        out, _ = self.to_out(out)
        return out

QKV Fusion in load_weights

When you fuse separate Q/K/V into QKVParallelLinear, map diffusers' separate weight names:

stacked_params_mapping = [
    ("to_qkv", "to_q", "q"),
    ("to_qkv", "to_k", "k"),
    ("to_qkv", "to_v", "v"),
]

def load_weights(self, weights):
    params = dict(self.named_parameters())
    loaded = set()
    for name, tensor in weights:
        for fused_name, orig_name, shard_id in stacked_params_mapping:
            if orig_name in name:
                name = name.replace(orig_name, fused_name)
                param = params[name]
                param.weight_loader(param, tensor, shard_id)
                loaded.add(name)
                break
        else:
            if name in params:
                param = params[name]
                if hasattr(param, "weight_loader"):
                    param.weight_loader(param, tensor)
                else:
                    default_weight_loader(param, tensor)
                loaded.add(name)
    return loaded

RMSNorm with TP

When RMSNorm sits between TP-sharded dimensions, use DistributedRMSNorm — it computes global RMS via all-reduce across TP ranks. See the Wan2.2 implementation for the pattern.

TP Constraints

• num_heads % tp_size == 0
• num_kv_heads % tp_size == 0
• Use self.to_qkv.num_heads (local per-GPU count), not total heads, for split sizes

Testing TP

python text_to_image.py --model Your-org/your-model \
  --tensor-parallel-size 2 --output "tp_test.png"

Verify: speedup, memory reduction proportional to TP size, quality matches single-GPU.

Reference implementations

Model	Path
Z-Image	vllm_omni/diffusion/models/z_image/z_image_transformer.py
FLUX	vllm_omni/diffusion/models/flux/flux_transformer.py
Qwen-Image	vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

---

Sequence Parallelism (SP / USP)

SP splits sequence tokens across GPUs using Ulysses (all-to-all) or Ring (P2P) communication. It is applied non-intrusively via the _sp_plan dict — no changes to forward() logic.

Approach 1: Non-Intrusive _sp_plan (Recommended)

The framework automatically registers hooks to shard inputs and gather outputs at nn.Module boundaries.

Step 1: Identify module boundaries

Find where tensors need sharding/gathering:

class MyTransformer(nn.Module):
    def __init__(self):
        self.patch_embed = PatchEmbed()    # Before blocks
        self.pos_embed = RoPE()            # RoPE may need splitting
        self.blocks = nn.ModuleList([...]) # Blocks process sharded x
        self.norm_out = LayerNorm()
        self.proj_out = Linear()           # Gather after this

    def forward(self, x):
        x = self.patch_embed(x)
        pos = self.pos_embed(x)
        for block in self.blocks:
            x = block(x, pos)
        x = self.norm_out(x)
        return self.proj_out(x)

Step 2: Handle inline operations

_sp_plan hooks only work at nn.Module boundaries. Inline ops like torch.cat() must be extracted into submodules:

# BAD: Inline — hooks can't intercept
unified = torch.cat([x, cap_feats], dim=1)

# GOOD: Extract into submodule
class UnifiedPrepare(nn.Module):
    def forward(self, x, cap_feats):
        return torch.cat([x, cap_feats], dim=1)

self.unified_prepare = UnifiedPrepare()
unified = self.unified_prepare(x, cap_feats)

Common cases: torch.cat(), pad_sequence(), tensor.reshape(), complex preprocessing.

Step 3: Write _sp_plan

Pattern 1: Shard at first block, gather at output (most common)

from vllm_omni.diffusion.distributed.sp_plan import (
    SequenceParallelInput, SequenceParallelOutput,
)

class StandardTransformer(nn.Module):
    _sp_plan = {
        "blocks.0": {
            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
        },
        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
    }

Pattern 2: Shard RoPE outputs separately

class TransformerWithRoPE(nn.Module):
    _sp_plan = {
        "rope": {
            0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
            1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
        },
        "blocks.0": {
            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
        },
        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
    }

Pattern 3: Dual-stream (shard image, replicate text)

class DualStreamTransformer(nn.Module):
    _sp_plan = {
        "rope_preparer": {
            2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
            3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
        },
        "transformer_blocks.0": {
            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
        },
        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
    }

API Reference

SequenceParallelInput:

Parameter	Type	Description
split_dim	int	Dimension to split (usually 1 for sequence)
expected_dims	int/None	Expected tensor rank for validation
split_output	bool	False: shard input params; True: shard output tensors
auto_pad	bool	Auto-pad if sequence not divisible by world_size

SequenceParallelOutput:

Parameter	Type	Description
gather_dim	int	Dimension to gather (usually 1 for sequence)
expected_dims	int/None	Expected tensor rank for validation

Module naming:

Key	Meaning
"blocks.0"	First element of ModuleList
"blocks.*"	All elements of ModuleList
"rope"	Named submodule

Dictionary value types:

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

Approach 2: Intrusive Modification (Complex Cases)

For dynamic sharding logic that can't 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)
    for block in self.blocks:
        hidden_states = block(hidden_states)
    if self.parallel_config.sequence_parallel_size > 1:
        hidden_states = sp_gather(hidden_states, dim=1)
    return hidden_states

Use intrusive modification as a last resort — _sp_plan is preferred for maintainability.

UAA Mode (Experimental)

ulysses_mode="advanced_uaa" handles arbitrary sequence lengths and head counts that aren't divisible by ulysses_degree. Uses variable all-to-all split sizes and temporary head padding.

Combining SP methods

Ulysses and Ring can be combined: ulysses_degree × ring_degree = total SP GPUs.

DiffusionParallelConfig(ulysses_degree=2, ring_degree=2)  # 4 GPUs total

Testing SP

# Offline
python text_to_image.py --model Your-model --ulysses-degree 2

# Online serving
vllm serve Your-model --omni --usp 2

Reference implementations

Model	Path
Qwen-Image	vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py
Wan2.2	vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py
Z-Image	vllm_omni/diffusion/models/z_image/z_image_transformer.py

---

CFG Parallelism

Distributes positive/negative Classifier-Free Guidance branches across 2 GPUs.

Implementation

Inherit CFGParallelMixin and implement diffuse():

from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin

class YourPipeline(nn.Module, CFGParallelMixin):
    def diffuse(self, latents, timesteps, prompt_embeds, negative_embeds,
                do_true_cfg, true_cfg_scale, **kwargs):
        for i, t in enumerate(timesteps):
            positive_kwargs = {
                "hidden_states": latents,
                "encoder_hidden_states": prompt_embeds,
                "timestep": t,
            }
            negative_kwargs = {
                "hidden_states": latents,
                "encoder_hidden_states": negative_embeds,
                "timestep": t,
            } if do_true_cfg else None

            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,
            )
            latents = self.scheduler_step_maybe_with_cfg(
                noise_pred, t, latents, do_true_cfg
            )
        return latents

Customization hooks

Method	Override when
predict_noise()	Non-standard transformer call (e.g., dual-transformer like Wan2.2)
cfg_normalize_function()	Custom normalization (e.g., LongCat with clamping)
combine_cfg_noise()	Multi-output models (e.g., video + audio: CFG on video, positive-only on audio)

Custom predict_noise (Wan2.2 — selects active transformer):

def predict_noise(self, current_model=None, **kwargs):
    if current_model is None:
        current_model = self.transformer
    return current_model(**kwargs)[0]

Custom combine_cfg_noise (multi-output):

def combine_cfg_noise(self, positive_pred, negative_pred, scale, normalize):
    video_pos, audio_pos = positive_pred
    video_neg, audio_neg = negative_pred
    video_combined = super().combine_cfg_noise(video_pos, video_neg, scale, normalize)
    return (video_combined, audio_pos)

Composite scheduler for multi-output

When each output has its own schedule:

class VideoAudioScheduler:
    def __init__(self, video_scheduler, audio_scheduler):
        self.video_scheduler = video_scheduler
        self.audio_scheduler = audio_scheduler

    def step(self, noise_pred, t, latents, return_dict=False, generator=None):
        video_out = self.video_scheduler.step(
            noise_pred[0], t[0], latents[0], return_dict=False, generator=generator
        )[0]
        audio_out = self.audio_scheduler.step(
            noise_pred[1], t[1], latents[1], return_dict=False, generator=generator
        )[0]
        return ((video_out, audio_out),)

Testing CFG Parallel

python text_to_image.py --model Your-model \
  --cfg-parallel-size 2 --cfg-scale 4.0 \
  --negative-prompt "ugly, unclear"

Constraint: guidance_scale > 1.0 and negative prompt must be provided.

Reference implementations

Model	Path
Qwen-Image	vllm_omni/diffusion/models/qwen_image/cfg_parallel.py
Wan2.2	vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
Mixin base	vllm_omni/diffusion/distributed/cfg_parallel.py

---

HSDP (Hybrid Sharded Data Parallel)

Shards model weights across GPUs using PyTorch FSDP2. Reduces per-GPU VRAM without changing computation.

Implementation

Add _hsdp_shard_conditions to the transformer class:

class YourTransformer(nn.Module):
    @staticmethod
    def _is_transformer_block(name: str, module) -> bool:
        return "blocks" in name and name.split(".")[-1].isdigit()

    _hsdp_shard_conditions = [_is_transformer_block]

For MoE models, add additional conditions:

class MoETransformer(nn.Module):
    @staticmethod
    def _is_transformer_block(name, module):
        return "blocks" in name and name.split(".")[-1].isdigit()

    @staticmethod
    def _is_moe_expert(name, module):
        return "experts" in name and name.split(".")[-1].isdigit()

    _hsdp_shard_conditions = [_is_transformer_block, _is_moe_expert]

A module is sharded if any condition returns True.

Constraints

• Cannot combine with Tensor Parallelism
• For standalone HSDP (no other parallelism), hsdp_shard_size must be specified explicitly
• Can combine with SP: HSDP reduces memory while SP distributes sequence

Testing HSDP

from vllm_omni.diffusion.data import DiffusionParallelConfig

parallel_config = DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=8)
omni = Omni(model="your-model", parallel_config=parallel_config)

Or CLI:

vllm serve Your-model --omni --use-hsdp

Verify: logs show "HSDP Inference: replicate_size=..., shard_size=..." and "Sharded N modules + root". Check VRAM reduction.

Reference implementations

Model	Path
Wan2.2	vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py
HSDP Core	vllm_omni/diffusion/distributed/hsdp.py

---

VAE Patch Parallelism

Shards VAE decode spatially across ranks using tiling:

python text_to_image.py --model Your-model --vae-patch-parallel-size 4

Auto-enables --vae-use-tiling. Uses DistributedAutoencoderKLWan or similar distributed VAE. Set vae_patch_parallel_size in DiffusionParallelConfig.

---

Combining Parallelism Methods

Common multi-GPU recipes:

# 4 GPUs: CFG (2) × Ulysses (2)
python text_to_image.py --model Qwen/Qwen-Image \
  --cfg-parallel-size 2 --ulysses-degree 2

# 8 GPUs: Ulysses (4) × Ring (2) + VAE patch (8)
python text_to_video.py --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
  --ulysses-degree 4 --ring-degree 2 --vae-patch-parallel-size 8

# 2 GPUs: HSDP + Ulysses (cannot combine HSDP with TP)
vllm serve Your-model --omni --use-hsdp --usp 2

Discovering Parallelism Support

Check which parallelism methods a model supports:

Check	How
Ulysses / Ring SP	Transformer defines _sp_plan. Search: grep -r '_sp_plan' vllm_omni/diffusion/models/
CFG Parallel	Pipeline inherits CFGParallelMixin. Search: grep -r 'CFGParallelMixin' vllm_omni/diffusion/models/
TP	Uses ColumnParallelLinear / QKVParallelLinear. Search: grep -r 'ParallelLinear|QKVParallel' vllm_omni/diffusion/models/<model>/
HSDP	Transformer defines _hsdp_shard_conditions. Search: grep -r '_hsdp_shard_conditions' vllm_omni/diffusion/models/

The canonical per-model support table is in docs/user_guide/diffusion/parallelism_acceleration.md.