This section describes how to add Sequence Parallel (SP) to a diffusion transformer model. We use the Qwen-Image transformer and Wan2.2 transformer as reference implementations.
**Terminology Note:** Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in the [diffusers library](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py). We use "Sequence Parallelism" to align with vLLM-Omni's terminology.
Diffusion transformers process long sequences of image patches or video frames. For high-resolution generation, these sequences can become very large. Enabling SP allows each GPU to process only a portion of the sequence, with attention mechanisms (Ulysses/Ring) handling cross-GPU communication transparently.
The `_sp_plan` mechanism allows SP **without modifying `forward()` logic**. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries.
**When to use:**
- Standard transformer architectures
- Tensor operations happen at `nn.Module` boundaries
- Predictable sharding/gathering patterns
**How it works:**
1. Declare `_sp_plan` dict in your transformer class
2. Framework automatically applies hooks when `sequence_parallel_size > 1`
3. Hooks shard/gather tensors at specified module boundaries
4. Attention layers handle cross-GPU communication internally
```python
classStandardTransformer(nn.Module):
_sp_plan={
# Shard hidden_states at first transformer block input
`StandardTransformer` has a transformer blocks list `self.blocks = nn.ModuleList([...])`, and a projection output layer `self.proj_out`. The `_sp_plan` above defines that when SP is enabled, sharding the input tensor to the first transformer block, and gathering the sharded tensor at the final output projection layer.
**Requirements:**
- Tensor operations that need sharding/gathering must happen at **`nn.Module` boundaries**
- Inline Python operations (e.g., `torch.cat`, `pad_sequence`) **cannot be hooked**
**Solution for inline operations:** Extract into a submodule (see Step 2 below).
### Step 1: Understand Module Boundaries
Identify where tensors need to be sharded or gathered in your model's forward pass:
```python
classMyTransformer(nn.Module):
def__init__(self):
self.patch_embed=PatchEmbed()# ← Boundary 1
self.pos_embed=RoPE()# ← Boundary 2
self.blocks=nn.ModuleList([...])# ← Boundary 3
self.norm_out=LayerNorm()
self.proj_out=Linear()# ← Boundary 4
defforward(self,x):
x=self.patch_embed(x)# ← Shard before this?
pos=self.pos_embed(x)# ← Shard RoPE outputs?
forblockinself.blocks:
x=block(x,pos)# ← Blocks process sharded x
x=self.norm_out(x)
output=self.proj_out(x)# ← Gather after this?
returnoutput
```
### Step 2: Handle Inline Operations
If your `forward()` contains inline tensor operations, **extract them into submodules**.
**Example: Z-Image concatenates image + text features inline**
For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls. Importantly, when taking this approach, be careful to ensure that you correctly manage the `_sp_shard_depth`; if the sequence parallel shard depth is 0, Ulysses will not be used.
After implementing Sequence Parallel support, thoroughly test your implementation to ensure correctness and performance across different configurations.
**Test Different `sp_size`:**
Test your model with various sequence parallel world sizes to verify correctness and identify optimal configurations:
**Solution:** In `SequenceParallelInput`, set `auto_pad=True` and add attention mask support.
> **Experimental Feature:** `auto_pad=True` is an experimental feature and may be changed in the future. We plan to improve this solution to involve minimal changes to model files. More details are [here](https://github.com/vllm-project/vllm-omni/issues/1324).
**Constraints of auto_pad:**
| Constraint | Description |
|------------|-------------|
| **Attention Backend Compatibility** | The attention backends must support `attention_mask`. Currently only `TORCH_SDPA` and `FLASH_ATTN` (default for diffusion models) are supported. |
| **Ring Attention Limitation** | Ring attention does not support `attention_mask`. Therefore, when using `auto_pad=True`, the combination of Ulysses + Ring attention is not feasible. |
1. Enable `auto_pad=True` for all sequence-dimension inputs in `_sp_plan`:
While `auto_pad` enables generation for irregular resolutions, be aware of potential quality impacts:
| Aspect | Impact |
|--------|--------|
| **Training Distribution** | Models perform best on aspect ratios seen during training (e.g., 1:1, 16:9, 4:3). Unusual ratios like 700x400 (1.75:1) may produce lower quality results. |
| **Padding Overhead** | Padded positions consume compute even when masked. For best efficiency, prefer resolutions divisible by `sp_size`. |
**Recommendations for users:**
- Use standard aspect ratios when possible (e.g., 768x432 for 16:9 instead of 700x400)
- Ensure post-patch dimensions are divisible by `sp_size` for optimal quality
- Test generation quality when using unusual resolutions
