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
This is the ideal approach for integrating sequence parallelism into new models, as it is easier to maintain and ensure compatibility with other types of acceleration.
`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**
**Pattern 3: Shard RoPE for Dual Stream Attention**
In some cases, different streams in attention may need to handle sequence parallelism differently. For example, we may want to shard the image embeddings, while replicating the text embeddings to correctly configure joint attention.
```python
classDualStreamTransformer(nn.Module):
"""
Dual-stream model where we need to replicate the text components, but shard
the image components to correctly handle sequence parallelism.
"""
_sp_plan={
# In this case, the rope_preparer returns a tuple of len 4, where the
# first 2 items correspond to the text, and the second 2 correspond to
NOTE: be careful to test adequately when refactoring classes that take this style of plan, as changing the order of the return values will break sequence parallelism.
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
