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.
`ulysses_mode="advanced_uaa"` enables the experimental UAA ("Ulysses Anything Attention") feature, which lets Ulysses attention handle arbitrary sequence lengths and arbitrary attention head counts. The same idea is also supported by [Cache-DiT](https://cache-dit.readthedocs.io/en/latest/user_guide/CONTEXT_PARALLEL/#uaa-ulysses-anything-attention).
Use it when plain Ulysses-SP would otherwise fail because:
- the local sequence shards are not evenly divisible after split hooks, or
- the attention head count is not divisible by `ulysses_degree`.
### Design Summary
1.**Strict mode stays unchanged.**
`ulysses_mode="strict"` keeps the original fast path and still requires divisible sequence/head shapes.
2.**UAA uses variable all-to-all split sizes for sequence shards.**
Before the Ulysses Q/K/V exchange, each rank all-gathers its local sequence length and uses those lengths as `all_to_all_single(..., output_split_sizes=seq_lens)`. This lets Ulysses gather the full sequence even when each rank started with a different local shard length.
3.**UAA pads heads only inside the Ulysses exchange.**
If `head_cnt % ulysses_degree != 0`, UAA pads the head dimension up to the next multiple of `ulysses_degree`, performs the forward/reverse all-to-all, then slices the temporary head padding away after the reverse exchange. The same rule is applied to joint attention tensors.
4.**Hybrid Ulysses + Ring is still shape-constrained.**
Ring attention expects every rank in a ring group to exchange exactly the same post-Ulysses sequence shape. UAA therefore validates those shapes before entering the ring path and raises a clear error if ring peers disagree on `S_global`.
5.**Tiny scalar gathers stay out of TorchDynamo tracing.**
`_all_gather_int()` is marked with `@torch.compiler.disable` so the scalar `item()` conversions used by UAA metadata collection do not get traced into `torch.compile`.
### UAA vs `auto_pad`
-`auto_pad=True` pads sequence tokens in `_sp_plan` and requires attention backends that support `attention_mask`.
-`advanced_uaa` does not depend on mask-based token padding inside Ulysses attention. It is therefore a better fit for non-divisible head counts and uneven Ulysses shard sizes.
-`auto_pad=True` remains incompatible with Ring attention because the ring backend does not consume `attention_mask`.
-`advanced_uaa` is still experimental and hybrid mode remains limited by Ring's equal-shape requirement.
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:
**Problem:** strict Ulysses sequence parallel requires divisible shapes. If the shard length is uneven, or if the model head count is not divisible by `ulysses_degree`, the strict path will raise an error.
1. Use `ulysses_mode="advanced_uaa"` for Ulysses-SP when you want the experimental uneven-shape path without relying on attention-mask padding.
2. If the model already uses `_sp_plan` token padding and the attention backend supports `attention_mask`, set `auto_pad=True` and add attention-mask plumbing.
> **Experimental Feature:** `ulysses_mode="advanced_uaa"` is experimental. It is intended to relax Ulysses divisibility constraints, but hybrid Ulysses + Ring still requires equal post-Ulysses sequence lengths inside each ring group.
> **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
