From c2888a8a36186ab5d3605bfb9b50adbb4c7c7db3 Mon Sep 17 00:00:00 2001
From: dongbo910220 <32610838+dongbo910220@users.noreply.github.com>
Date: Wed, 25 Mar 2026 15:46:15 +0800
Subject: [PATCH] [Docs] refine UAA documentation (#2148)

Signed-off-by: dongbo910220 <1275604947@qq.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
---
 docs/design/feature/sequence_parallel.md              | 2 +-
 docs/user_guide/diffusion/parallelism_acceleration.md | 9 ++++-----
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/docs/design/feature/sequence_parallel.md b/docs/design/feature/sequence_parallel.md
index c70a148b9..d0328bcf6 100644
--- a/docs/design/feature/sequence_parallel.md
+++ b/docs/design/feature/sequence_parallel.md
@@ -49,7 +49,7 @@ from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
 
 ## UAA Mode (Experimental)
 
-`ulysses_mode="advanced_uaa"` enables the experimental UAA ("Ulysses Anything Attention") path that was introduced to address the uneven-shape cases discussed in [issue #1324](https://github.com/vllm-project/vllm-omni/issues/1324).
+`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:
 
diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md
index 0d6903a6a..b6c8aed59 100644
--- a/docs/user_guide/diffusion/parallelism_acceleration.md
+++ b/docs/user_guide/diffusion/parallelism_acceleration.md
@@ -135,7 +135,7 @@ 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.
+    `ulysses_mode="advanced_uaa"` is an experimental extension to Ulysses-SP. It lets Ulysses attention handle arbitrary sequence lengths and arbitrary attention head counts 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.
 
@@ -161,7 +161,7 @@ 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:
+To enable the experimental UAA mode explicitly, use a model/configuration that actually requires it. For example, `Tongyi-MAI/Z-Image-Turbo` has 30 attention heads, so `ulysses_degree=4` requires UAA because 30 is not divisible by 4:
 
 ```python
 omni = Omni(
@@ -181,7 +181,7 @@ You can enable Ulysses-SP in online serving for diffusion models via `--usp`:
 # Text-to-image (requires >= 2 GPUs)
 vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2
 
-# Experimental UAA mode
+# Experimental UAA mode for a model with 30 attention heads
 vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 --usp 4 --ulysses-mode advanced_uaa
 ```
 
@@ -279,7 +279,6 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(
         ulysses_degree=2,
         ring_degree=2,
-        ulysses_mode="advanced_uaa",
     )
 )
 
@@ -293,7 +292,7 @@ outputs = omni.generate(
 
 ```bash
 # Text-to-image (requires >= 4 GPUs)
-vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2 --ulysses-mode advanced_uaa
+vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2
 ```
 
 ##### Benchmarks