anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Feature]: add Ulysses advanced_uaa mode (#1379)

Browse Source 
Signed-off-by: dongbo910220 <1275604947@qq.com>
Signed-off-by: dongbo910220 <32610838+dongbo910220@users.noreply.github.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
This commit is contained in:
dongbo910220andDidan Deng authored and GitHub committed 2026-03-24 14:55:54 +08:00
1 parent 4cfa9bf843
commit 50bd1961bf
17 files changed
+1064 -13
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/user_guide/diffusion/parallelism_acceleration.md
+29 -2
View File
@@ -134,6 +134,11 @@ VAE patch parallelism automatically selects between two internal decode methods
#### Ulysses-SP
!!! note "Experimental UAA mode"
`ulysses_mode="advanced_uaa"` is an experimental extension to Ulysses-SP. It allows Ulysses attention to handle non-divisible attention head counts and uneven per-rank sequence shards without relying on `attention_mask`-based token padding.
In hybrid Ulysses + Ring mode, Ring still requires every rank in the same ring group to observe the same post-Ulysses sequence length. If that condition is not met, vLLM-Omni raises a validation error instead of entering the ring kernel with inconsistent shapes.
##### Offline Inference
An example of offline inference script using [Ulysses-SP](https://arxiv.org/pdf/2309.14509) is shown below:
@@ -156,6 +161,18 @@ outputs = omni.generate(
See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
To enable the experimental UAA mode explicitly:
```python
omni = Omni(
model="Tongyi-MAI/Z-Image-Turbo",
parallel_config=DiffusionParallelConfig(
ulysses_degree=4,
ulysses_mode="advanced_uaa",
),
)
```
##### Online Serving
You can enable Ulysses-SP in online serving for diffusion models via `--usp`:
@@ -163,6 +180,9 @@ You can enable Ulysses-SP in online serving for diffusion models via `--usp`:
```bash
# Text-to-image (requires >= 2 GPUs)
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2
# Experimental UAA mode
vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 --usp 4 --ulysses-mode advanced_uaa
```
##### Benchmarks
@@ -243,6 +263,9 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m
You can combine both Ulysses-SP and Ring-Attention for larger scale parallelism. The total sequence parallel size equals `ulysses_degree × ring_degree`.
!!! note "Experimental UAA in hybrid mode"
`ulysses_mode="advanced_uaa"` can also be used with hybrid Ulysses + Ring, but this does not remove Ring's shape requirement. Every rank in the same ring group must still have the same post-Ulysses sequence length.
##### Offline Inference
```python
@@ -253,7 +276,11 @@ 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)
parallel_config=DiffusionParallelConfig(
ulysses_degree=2,
ring_degree=2,
ulysses_mode="advanced_uaa",
)
)
outputs = omni.generate(
@@ -266,7 +293,7 @@ outputs = omni.generate(
```bash
# Text-to-image (requires >= 4 GPUs)
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2 --ulysses-mode advanced_uaa
```
##### Benchmarks