diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index b606648c7..718f0c3e0 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -112,6 +112,7 @@ steps: depends_on: upload-merge-pipeline commands: - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py + - pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py agents: queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU plugins: diff --git a/docs/design/feature/sequence_parallel.md b/docs/design/feature/sequence_parallel.md index 3b899ffe3..c70a148b9 100644 --- a/docs/design/feature/sequence_parallel.md +++ b/docs/design/feature/sequence_parallel.md @@ -7,6 +7,7 @@ This section describes how to add Sequence Parallel (SP) to a diffusion transfor ## Table of Contents - [Overview](#overview) +- [UAA Mode (Experimental)](#uaa-mode-experimental) - [Approach 1: Non-Intrusive `_sp_plan` (Recommended)](#approach-1-non-intrusive-_sp_plan-recommended) - [Approach 2: Intrusive Modification (For Complex Cases)](#approach-2-intrusive-modification-for-complex-cases) - [Testing](#testing) @@ -46,6 +47,41 @@ 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). + +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. + +--- + ## Approach 1: Non-Intrusive `_sp_plan` (Recommended) The `_sp_plan` mechanism allows SP **without modifying `forward()` logic**. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries. @@ -364,9 +400,14 @@ _sp_plan = { - **Sequence Length not divisible by sp_size:** -**Problem:** `SequenceParallelInput(auto_pad=False)` - auto_pad should be True to enable automatic sequence padding. +**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. -**Solution:** In `SequenceParallelInput`, set `auto_pad=True` and add attention mask support. +**Solutions:** + +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). diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index dbae38e86..0d6903a6a 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -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 diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index 7019bd47b..61013068a 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -84,6 +84,9 @@ python text_to_image.py \ | `--vae-use-slicing` | flag | off | Enable VAE slicing for memory optimization | | `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization | | `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel | +| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree for multi-GPU inference | +| `--ring-degree` | int | `1` | Ring sequence parallel degree for hybrid Ulysses + Ring inference | +| `--ulysses-mode` | str | `"strict"` | Ulysses SP mode: `"strict"` or `"advanced_uaa"` | | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models | | `--lora-path` | str | — | Path to PEFT LoRA adapter folder | | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights | @@ -120,6 +123,20 @@ python text_to_image.py \ `Tongyi-MAI/Z-Image-Turbo` is a distilled version of Z-Image. Distilled diffusion models usually require less number of inference steps (4~9), and Classifier-Free Guidance (CFG) is usually NOT applied. Similar distilled models are `black-forest-labs/FLUX.2-klein-4B` and `black-forest-labs/FLUX.2-klein-9B`. +Hybrid Ulysses + Ring example (requires 4 GPUs): + +```bash +python text_to_image.py \ + --model Tongyi-MAI/Z-Image-Turbo \ + --prompt "a cup of coffee on the table" \ + --ulysses-degree 2 \ + --ring-degree 2 \ + --ulysses-mode advanced_uaa \ + --height 1024 \ + --width 1024 \ + --output outputs/coffee_hybrid.png +``` + ### NextStep Models NextStep-1.1 supports extra arguments for dual-level CFG control: diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 81189c9d6..8fd8cad90 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -109,6 +109,13 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of GPUs used for ulysses sequence parallelism.", ) + parser.add_argument( + "--ulysses-mode", + type=str, + default="strict", + choices=["strict", "advanced_uaa"], + help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA).", + ) parser.add_argument( "--ring-degree", type=int, @@ -272,6 +279,7 @@ def main(): parallel_config = DiffusionParallelConfig( ulysses_degree=args.ulysses_degree, ring_degree=args.ring_degree, + ulysses_mode=args.ulysses_mode, cfg_parallel_size=args.cfg_parallel_size, tensor_parallel_size=args.tensor_parallel_size, vae_patch_parallel_size=args.vae_patch_parallel_size, @@ -344,8 +352,10 @@ def main(): print(f" Ignored layers: {ignored_layers}") print( f" Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, " - f"ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, " - f"vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}." + f"ulysses_degree={args.ulysses_degree}, ulysses_mode={args.ulysses_mode}, " + f"ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, " + f"vae_patch_parallel_size={args.vae_patch_parallel_size}, " + f"enable_expert_parallel={args.enable_expert_parallel}." ) print(f" CPU offload: {args.enable_cpu_offload}") print(f" Image size: {args.width}x{args.height}") diff --git a/tests/diffusion/attention/test_ulysses_uaa.py b/tests/diffusion/attention/test_ulysses_uaa.py new file mode 100644 index 000000000..67907d7ef --- /dev/null +++ b/tests/diffusion/attention/test_ulysses_uaa.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import os +import socket +import tempfile + +import numpy as np +import pytest +import torch + +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_distributed_env, + init_distributed_environment, + initialize_model_parallel, +) +from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context +from vllm_omni.platforms import current_omni_platform + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return int(s.getsockname()[1]) + + +def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None: + os.environ["RANK"] = str(rank) + os.environ["LOCAL_RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(master_port) + + +def _run_attention_case( + local_rank: int, + world_size: int, + master_port: int, + input_file: str, + output_file: str, + num_heads: int, + head_size: int, + ulysses_degree: int, + ulysses_mode: str, + ring_degree: int = 1, + split_sizes: list[int] | None = None, + sdp_kernel_mode: str = "math", +) -> None: + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) + + _set_dist_env(rank=local_rank, world_size=world_size, master_port=master_port) + os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA" + + init_distributed_environment(world_size=world_size, rank=local_rank) + initialize_model_parallel( + data_parallel_size=1, + cfg_parallel_size=1, + sequence_parallel_size=world_size, + ulysses_degree=ulysses_degree, + ring_degree=ring_degree, + tensor_parallel_size=1, + pipeline_parallel_size=1, + ) + + parallel_config = DiffusionParallelConfig( + pipeline_parallel_size=1, + data_parallel_size=1, + tensor_parallel_size=1, + sequence_parallel_size=world_size, + ulysses_degree=ulysses_degree, + ring_degree=ring_degree, + cfg_parallel_size=1, + ulysses_mode=ulysses_mode, + ) + od_config = OmniDiffusionConfig(model="test", dtype=torch.float32, parallel_config=parallel_config) + + with set_forward_context(omni_diffusion_config=od_config): + attn = Attention( + num_heads=num_heads, + head_size=head_size, + causal=False, + softmax_scale=1.0 / (head_size**0.5), + ).to(device=device, dtype=torch.float32) + + with np.load(input_file, allow_pickle=False) as payload: + q_full = torch.from_numpy(payload["q"]).to(device=device) + k_full = torch.from_numpy(payload["k"]).to(device=device) + v_full = torch.from_numpy(payload["v"]).to(device=device) + + if world_size == 1: + q, k, v = q_full, k_full, v_full + else: + if split_sizes is None: + # NOTE: torch.chunk may return fewer than `world_size` chunks for some + # uneven lengths (e.g. seq_len=9, world_size=4 -> 3 chunks of len 3). + # We need an exact `world_size`-way split to simulate uneven SP shards. + q = torch.tensor_split(q_full, world_size, dim=1)[local_rank].contiguous() + k = torch.tensor_split(k_full, world_size, dim=1)[local_rank].contiguous() + v = torch.tensor_split(v_full, world_size, dim=1)[local_rank].contiguous() + else: + if len(split_sizes) != world_size: + raise ValueError(f"split_sizes length ({len(split_sizes)}) must equal world_size ({world_size}).") + if sum(int(x) for x in split_sizes) != q_full.shape[1]: + raise ValueError( + "split_sizes must sum to full seq_len " + f"(got sum={sum(int(x) for x in split_sizes)}, seq_len={q_full.shape[1]})." + ) + q = torch.split(q_full, split_sizes, dim=1)[local_rank].contiguous() + k = torch.split(k_full, split_sizes, dim=1)[local_rank].contiguous() + v = torch.split(v_full, split_sizes, dim=1)[local_rank].contiguous() + # The Attention layer only enables SP communication when ForwardContext.sp_active is True. + # In production this is managed by SequenceParallelSplitHook/GatherHook, but here we + # shard manually for testing. + get_forward_context()._sp_shard_depth = 1 + + if sdp_kernel_mode == "math": + sdp_ctx = torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True) + elif sdp_kernel_mode == "mem_efficient": + sdp_ctx = torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=True, enable_math=False) + else: + raise ValueError(f"Invalid sdp_kernel_mode: {sdp_kernel_mode!r}") + + with sdp_ctx: + out_local = attn(q, k, v).contiguous() + + if world_size == 1: + out_full = out_local + else: + local_len = torch.tensor([out_local.shape[1]], device=device, dtype=torch.int64) + gathered_lens = [torch.empty_like(local_len) for _ in range(world_size)] + torch.distributed.all_gather(gathered_lens, local_len) + lens = [int(t.item()) for t in gathered_lens] + max_len = max(lens) + + if out_local.shape[1] < max_len: + pad = max_len - out_local.shape[1] + out_local = torch.nn.functional.pad(out_local, (0, 0, 0, 0, 0, pad)).contiguous() + + gathered = [torch.empty_like(out_local) for _ in range(world_size)] + torch.distributed.all_gather(gathered, out_local) + if local_rank == 0: + out_full = torch.cat([t[:, : lens[i]].contiguous() for i, t in enumerate(gathered)], dim=1) + else: + out_full = None + + if local_rank == 0: + np.save(output_file, out_full.detach().cpu().numpy()) + + destroy_distributed_env() + + +@pytest.mark.parametrize( + "sp_world_size,seq_len,num_heads", + [ + (2, 6, 3), # head_cnt not divisible by P=2 + (2, 5, 4), # seq_len not divisible by P=2 + (4, 9, 30), # Z-Image-like: head_cnt not divisible by P=4 + (4, 10, 8), # seq_len not divisible by P=4 + ], +) +def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_heads: int) -> None: + if current_omni_platform.get_device_count() < sp_world_size: + pytest.skip(f"Test requires {sp_world_size} GPUs") + + batch_size = 2 + head_size = 8 + + base_port = _find_free_port() + sp_port = _find_free_port() + while sp_port == base_port: + sp_port = _find_free_port() + + with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in: + input_file = f_in.name + with tempfile.NamedTemporaryFile(delete=False, suffix=".npy") as f_base: + baseline_file = f_base.name + with tempfile.NamedTemporaryFile(delete=False, suffix=".npy") as f_sp: + sp_file = f_sp.name + + try: + torch.manual_seed(0) + q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32) + k = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32) + v = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32) + np.savez(input_file, q=q.numpy(), k=k.numpy(), v=v.numpy()) + + # Baseline (no SP) + torch.multiprocessing.spawn( + _run_attention_case, + args=(1, base_port, input_file, baseline_file, num_heads, head_size, 1, "strict"), + nprocs=1, + ) + + # SP (Ulysses-P with UAA) + torch.multiprocessing.spawn( + _run_attention_case, + args=( + sp_world_size, + sp_port, + input_file, + sp_file, + num_heads, + head_size, + sp_world_size, + "advanced_uaa", + ), + nprocs=sp_world_size, + ) + + baseline = np.load(baseline_file, allow_pickle=False) + sp = np.load(sp_file, allow_pickle=False) + + baseline_t = torch.from_numpy(baseline) + sp_t = torch.from_numpy(sp) + assert baseline_t.shape == sp_t.shape + torch.testing.assert_close(sp_t, baseline_t, atol=1e-5, rtol=1e-5) + finally: + for path in (input_file, baseline_file, sp_file): + try: + os.remove(path) + except OSError: + pass + + +def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None: + sp_world_size = 4 + ulysses_degree = 2 + ring_degree = 2 + + if current_omni_platform.get_device_count() < sp_world_size: + pytest.skip(f"Test requires {sp_world_size} GPUs") + + batch_size = 2 + head_size = 8 + seq_len = 10 + num_heads = 3 # head_cnt not divisible by ulysses_degree=2 -> triggers head padding + + # Ensure ring ranks see equal post-Ulysses seq_len: + # rank0/1 -> 3+2=5, rank2/3 -> 3+2=5 + split_sizes = [3, 2, 3, 2] + + base_port = _find_free_port() + sp_port = _find_free_port() + while sp_port == base_port: + sp_port = _find_free_port() + + with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in: + input_file = f_in.name + with tempfile.NamedTemporaryFile(delete=False, suffix=".npy") as f_base: + baseline_file = f_base.name + with tempfile.NamedTemporaryFile(delete=False, suffix=".npy") as f_sp: + sp_file = f_sp.name + + try: + torch.manual_seed(0) + q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32) + k = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32) + v = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32) + np.savez(input_file, q=q.numpy(), k=k.numpy(), v=v.numpy()) + + # Baseline (no SP) + torch.multiprocessing.spawn( + _run_attention_case, + args=(1, base_port, input_file, baseline_file, num_heads, head_size, 1, "strict", 1, None, "mem_efficient"), + nprocs=1, + ) + + # Hybrid SP: Ulysses (P=2) + Ring (P=2) with advanced_uaa + torch.multiprocessing.spawn( + _run_attention_case, + args=( + sp_world_size, + sp_port, + input_file, + sp_file, + num_heads, + head_size, + ulysses_degree, + "advanced_uaa", + ring_degree, + split_sizes, + "mem_efficient", + ), + nprocs=sp_world_size, + ) + + baseline = np.load(baseline_file, allow_pickle=False) + sp = np.load(sp_file, allow_pickle=False) + + baseline_t = torch.from_numpy(baseline) + sp_t = torch.from_numpy(sp) + assert baseline_t.shape == sp_t.shape + # Hybrid (Ulysses+Ring) typically has slightly larger numerical differences + # than pure Ulysses due to different communication/reduction order and + # the SDPA kernel path used by Ring attention. Use a looser tolerance to + # keep the test stable across GPUs/kernels while still catching regressions. + torch.testing.assert_close(sp_t, baseline_t, atol=5e-4, rtol=5e-4) + finally: + for path in (input_file, baseline_file, sp_file): + try: + os.remove(path) + except OSError: + pass diff --git a/tests/diffusion/distributed/test_sp_plan_hooks.py b/tests/diffusion/distributed/test_sp_plan_hooks.py index 72ec9aaf2..f95dbfe2b 100644 --- a/tests/diffusion/distributed/test_sp_plan_hooks.py +++ b/tests/diffusion/distributed/test_sp_plan_hooks.py @@ -988,6 +988,46 @@ class TestDimensionValidation: assert result.shape == tensor_4d.shape +@pytest.mark.cpu +class TestStrictModeSplitValidation: + """Test strict mode divisibility validation in SequenceParallelSplitHook.""" + + def test_strict_mode_raises_on_non_divisible_seq_len(self): + """Strict Ulysses-SP should fail fast when seq_len is not divisible by sp_size.""" + from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig + from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig + from vllm_omni.diffusion.forward_context import set_forward_context + from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook + + class DummyModule(nn.Module): + def forward(self, x): + return x + + metadata = {"x": SequenceParallelInput(split_dim=1, expected_dims=3)} + config = SequenceParallelConfig(ulysses_degree=2, ring_degree=1) + hook = SequenceParallelSplitHook(metadata, config) + hook.initialize_hook(DummyModule()) + + # seq_len=5 is not divisible by sp_size=2 + x = torch.randn(1, 5, 8) + + parallel_config = DiffusionParallelConfig( + pipeline_parallel_size=1, + data_parallel_size=1, + tensor_parallel_size=1, + sequence_parallel_size=2, + ulysses_degree=2, + ring_degree=1, + cfg_parallel_size=1, + ulysses_mode="strict", + ) + od_config = OmniDiffusionConfig(model="test", dtype=torch.float32, parallel_config=parallel_config) + + with set_forward_context(omni_diffusion_config=od_config): + with pytest.raises(ValueError, match=r"strict mode.*sequence_parallel_size"): + hook._prepare_sp_input(x, metadata["x"], (), {}) + + @pytest.mark.cpu class TestSequenceParallelConfig: """Test SequenceParallelConfig dataclass.""" diff --git a/tests/diffusion/distributed/test_ulysses_uaa_perf.py b/tests/diffusion/distributed/test_ulysses_uaa_perf.py new file mode 100644 index 000000000..c8b07ba15 --- /dev/null +++ b/tests/diffusion/distributed/test_ulysses_uaa_perf.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Perf smoke tests for Ulysses advanced_uaa communication overhead. + +This test is intended for CI monitoring only: +- Print per-iteration timings and ratio vs strict Ulysses all-to-all. +- Use a loose sanity bound to catch gross regressions without flakiness. +""" + +from __future__ import annotations + +import os +import socket +from dataclasses import dataclass + +import pytest +import torch +import torch.distributed as dist + +from vllm_omni.diffusion.attention.parallel.ulysses import ( + _all_gather_int, + _ulysses_all_to_all_any_o, + _ulysses_all_to_all_any_qkv, +) +from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D +from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_distributed_env, + get_sp_group, + init_distributed_environment, + initialize_model_parallel, +) +from vllm_omni.platforms import current_omni_platform + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return int(s.getsockname()[1]) + + +def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None: + os.environ["RANK"] = str(rank) + os.environ["LOCAL_RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(master_port) + + +def _max_all_reduce(pg: dist.ProcessGroup, value: float, *, device: torch.device) -> float: + t = torch.tensor([value], device=device, dtype=torch.float32) + dist.all_reduce(t, op=dist.ReduceOp.MAX, group=pg) + return float(t.item()) + + +@dataclass(frozen=True, slots=True) +class _PerfCase: + ulysses_degree: int + ring_degree: int + + @property + def world_size(self) -> int: + return int(self.ulysses_degree * self.ring_degree) + + +PERF_CASES: list[_PerfCase] = [ + _PerfCase(ulysses_degree=4, ring_degree=1), + _PerfCase(ulysses_degree=2, ring_degree=2), +] + + +@pytest.mark.parametrize("case", PERF_CASES) +def test_ulysses_advanced_uaa_comm_overhead(case: _PerfCase) -> None: + available_gpus = current_omni_platform.get_device_count() + if available_gpus < case.world_size: + pytest.skip(f"Requires {case.world_size} GPUs, got {available_gpus}") + + master_port = _find_free_port() + torch.multiprocessing.spawn( + _perf_worker, + args=(case.world_size, master_port, case.ulysses_degree, case.ring_degree), + nprocs=case.world_size, + ) + + +def _perf_worker(local_rank: int, world_size: int, master_port: int, ulysses_degree: int, ring_degree: int) -> None: + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) + + _set_dist_env(rank=local_rank, world_size=world_size, master_port=master_port) + + try: + init_distributed_environment(world_size=world_size, rank=local_rank) + initialize_model_parallel( + data_parallel_size=1, + cfg_parallel_size=1, + sequence_parallel_size=world_size, + ulysses_degree=ulysses_degree, + ring_degree=ring_degree, + tensor_parallel_size=1, + pipeline_parallel_size=1, + ) + + sp_group = get_sp_group() + ulysses_pg = sp_group.ulysses_group + ring_pg = sp_group.ring_group + + ulysses_world_size = dist.get_world_size(ulysses_pg) + ring_world_size = dist.get_world_size(ring_pg) + + # A moderate tensor size to reduce timing noise while staying fast in CI. + bsz = 1 + s_local = 256 + head_cnt = 32 # divisible by ulysses_degree in both cases above + head_dim = 128 + dtype = torch.float16 + use_sync = False + + torch.manual_seed(1234 + local_rank) + q = torch.randn(bsz, s_local, head_cnt, head_dim, device=device, dtype=dtype) + k = torch.randn_like(q) + v = torch.randn_like(q) + + warmup_iters = 10 + iters = 200 + + def strict_comm_step(q_in: torch.Tensor, k_in: torch.Tensor, v_in: torch.Tensor) -> torch.Tensor: + q_out = SeqAllToAll4D.apply(ulysses_pg, q_in, 2, 1, use_sync) + SeqAllToAll4D.apply(ulysses_pg, k_in, 2, 1, use_sync) + SeqAllToAll4D.apply(ulysses_pg, v_in, 2, 1, use_sync) + o_in = q_out # dummy + return SeqAllToAll4D.apply(ulysses_pg, o_in, 1, 2, use_sync) + + def uaa_comm_step(q_in: torch.Tensor, k_in: torch.Tensor, v_in: torch.Tensor) -> torch.Tensor: + seq_lens = _all_gather_int(ulysses_pg, int(q_in.shape[1]), device=q_in.device) + s_global = int(sum(seq_lens)) + if ring_world_size > 1: + ring_s_globals = _all_gather_int(ring_pg, s_global, device=q_in.device) + if len(set(ring_s_globals)) != 1: + raise RuntimeError(f"Unexpected hybrid ring post-Ulysses seq_len mismatch: {ring_s_globals}.") + + q_out, orig_head_cnt = _ulysses_all_to_all_any_qkv(ulysses_pg, q_in, seq_lens=seq_lens, use_sync=use_sync) + _ulysses_all_to_all_any_qkv(ulysses_pg, k_in, seq_lens=seq_lens, use_sync=use_sync) + _ulysses_all_to_all_any_qkv(ulysses_pg, v_in, seq_lens=seq_lens, use_sync=use_sync) + o_in = q_out # dummy + return _ulysses_all_to_all_any_o( + ulysses_pg, + o_in, + seq_lens=seq_lens, + local_seq_len=int(s_local), + orig_head_cnt=int(orig_head_cnt), + use_sync=use_sync, + ) + + with torch.no_grad(): + # Warmup (strict) + for _ in range(warmup_iters): + _ = strict_comm_step(q, k, v) + current_omni_platform.synchronize() + + # Timed (strict) + t0 = torch.cuda.Event(enable_timing=True) + t1 = torch.cuda.Event(enable_timing=True) + t0.record() + for _ in range(iters): + _ = strict_comm_step(q, k, v) + t1.record() + current_omni_platform.synchronize() + strict_ms = float(t0.elapsed_time(t1)) / float(iters) + + # Warmup (UAA) + for _ in range(warmup_iters): + _ = uaa_comm_step(q, k, v) + current_omni_platform.synchronize() + + # Timed (UAA) + u0 = torch.cuda.Event(enable_timing=True) + u1 = torch.cuda.Event(enable_timing=True) + u0.record() + for _ in range(iters): + _ = uaa_comm_step(q, k, v) + u1.record() + current_omni_platform.synchronize() + uaa_ms = float(u0.elapsed_time(u1)) / float(iters) + + # Reduce across ranks (use worst-rank to be conservative). + strict_ms_max = _max_all_reduce(dist.group.WORLD, strict_ms, device=device) + uaa_ms_max = _max_all_reduce(dist.group.WORLD, uaa_ms, device=device) + ratio = (uaa_ms_max / strict_ms_max) if strict_ms_max > 0 else float("inf") + + # Approx bytes moved per iteration per-rank (send+recv) for 4x all-to-all. + elem_size = torch.tensor([], dtype=dtype).element_size() + per_a2a_bytes = int(q.numel()) * elem_size * 2 + comm_bytes = int(4 * per_a2a_bytes) + strict_gbps = (comm_bytes / (strict_ms_max / 1000.0)) / 1e9 if strict_ms_max > 0 else 0.0 + uaa_gbps = (comm_bytes / (uaa_ms_max / 1000.0)) / 1e9 if uaa_ms_max > 0 else 0.0 + + if dist.get_rank() == 0: + payload = { + "name": "ulysses_uaa_comm_perf", + "world_size": int(world_size), + "ulysses_degree": int(ulysses_degree), + "ring_degree": int(ring_degree), + "ulysses_world_size": int(ulysses_world_size), + "ring_world_size": int(ring_world_size), + "shape": { + "bsz": int(bsz), + "s_local": int(s_local), + "head_cnt": int(head_cnt), + "head_dim": int(head_dim), + }, + "dtype": str(dtype), + "iters": int(iters), + "strict_ms_per_iter_max": float(strict_ms_max), + "uaa_ms_per_iter_max": float(uaa_ms_max), + "uaa_over_strict_ratio": float(ratio), + "comm_bytes_per_iter_per_rank": int(comm_bytes), + "strict_effective_gbps_per_rank": float(strict_gbps), + "uaa_effective_gbps_per_rank": float(uaa_gbps), + } + print(f"UAA_COMM_PERF_JSON={payload}") + + # Loose bound: we only want to catch severe regressions. The printed JSON + # payload is used for monitoring smaller changes, while this threshold is + # intentionally generous to avoid flakiness across GPU types/drivers. + max_uaa_ms_per_iter = 10.0 + assert uaa_ms_max < max_uaa_ms_per_iter, ( + f"UAA comm too slow: uaa={uaa_ms_max:.3f}ms/iter, strict={strict_ms_max:.3f}ms/iter " + f"(cap={max_uaa_ms_per_iter:.1f}ms/iter)." + ) + assert ratio < 3.0, f"UAA comm overhead too high: ratio={ratio:.3f}x (strict={strict_ms_max:.3f}ms)." + finally: + destroy_distributed_env() diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 06939a7d1..ca5624f2d 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -2,8 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.entrypoints.cli.serve import OmniServeCommand, _create_default_diffusion_stage_cfg pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -51,3 +53,43 @@ def test_default_stage_devices_from_sequence_parallel(): )[0] assert stage_cfg["runtime"]["devices"] == "0,1,2,3" + + +def test_default_stage_config_propagates_ulysses_mode(): + """Ensure UAA mode survives default diffusion-stage creation.""" + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + "ulysses_degree": 4, + "ulysses_mode": "advanced_uaa", + } + )[0] + + parallel_config = stage_cfg["engine_args"]["parallel_config"] + assert parallel_config.ulysses_degree == 4 + assert parallel_config.ulysses_mode == "advanced_uaa" + + +def test_serve_cli_accepts_ulysses_mode(): + """Ensure diffusion serve CLI exposes ulysses_mode and wires it to parallel_config.""" + parser = FlexibleArgumentParser() + subparsers = parser.add_subparsers(dest="command") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args( + [ + "serve", + "Qwen/Qwen-Image", + "--omni", + "--usp", + "4", + "--ulysses-mode", + "advanced_uaa", + ] + ) + + stage_cfg = _create_default_diffusion_stage_cfg(args)[0] + parallel_config = stage_cfg["engine_args"]["parallel_config"] + + assert args.ulysses_mode == "advanced_uaa" + assert parallel_config.ulysses_degree == 4 + assert parallel_config.ulysses_mode == "advanced_uaa" diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py b/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py index ad61f9a20..51cf8bbc3 100644 --- a/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py +++ b/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py @@ -62,6 +62,21 @@ def pytorch_attn_forward( if op_type == "flash" and q.dtype == torch.float32: op_type = "efficient" + # Volta (V100, sm70) does not have native bf16 support for the efficient + # SDPA kernel. When running Ring attention in bf16 on such GPUs, the + # efficient kernel can fail with "cutlassF: no kernel found to launch". + # + # We keep the overall model in bf16 for numerical stability, but cast + # the attention inputs to fp16 for the SDPA call. The output is cast + # back to bf16 to match the caller's expected dtype. + orig_dtype = q.dtype + if q.is_cuda and q.dtype == torch.bfloat16 and op_type == "efficient": + major_cc, _minor_cc = torch.cuda.get_device_capability(q.device) + if major_cc < 8: + q = q.to(torch.float16) + k = k.to(torch.float16) + v = v.to(torch.float16) + q = q.transpose(1, 2) k = k.transpose(1, 2) v = v.transpose(1, 2) @@ -90,7 +105,13 @@ def pytorch_attn_forward( raise ValueError(f"Invalid op_type: {op_type}") out = out.transpose(1, 2) - lse = lse.to(q.dtype) + # Keep LSE in fp32 for numerical stability when accumulating across ring + # steps (update_out_and_lse uses sigmoid/logsigmoid on LSE diffs). + # Casting LSE down to fp16/bf16 can introduce NaNs on some GPUs/shapes. + lse = lse.to(torch.float32) + + if out.dtype != orig_dtype: + out = out.to(orig_dtype) return out, lse diff --git a/vllm_omni/diffusion/attention/parallel/ulysses.py b/vllm_omni/diffusion/attention/parallel/ulysses.py index 66914102e..db6e1123b 100644 --- a/vllm_omni/diffusion/attention/parallel/ulysses.py +++ b/vllm_omni/diffusion/attention/parallel/ulysses.py @@ -7,11 +7,143 @@ from dataclasses import dataclass import torch import torch.distributed as dist +import torch.nn.functional as F from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.parallel.base import ParallelAttentionContext from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGroupCoordinator +from vllm_omni.diffusion.forward_context import get_ulysses_mode + + +def _ceil_div(n: int, d: int) -> int: + return (n + d - 1) // d + + +def _positive_divisors(n: int) -> list[int]: + if n <= 0: + return [] + divs = set() + i = 1 + while i * i <= n: + if n % i == 0: + divs.add(i) + divs.add(n // i) + i += 1 + return sorted(divs) + + +@torch.compiler.disable +def _all_gather_int(pg: dist.ProcessGroup, value: int, *, device: torch.device) -> list[int]: + """All-gather a scalar int across pg. + + Note: we use a device tensor so this works for NCCL subgroups (e.g. Ulysses/Ring). + """ + world_size = dist.get_world_size(pg) + if world_size == 1: + return [int(value)] + + t = torch.tensor([int(value)], dtype=torch.int64, device=device) + gathered = [torch.empty_like(t) for _ in range(world_size)] + dist.all_gather(gathered, t, group=pg) + return [int(x.item()) for x in gathered] + + +def _ulysses_all_to_all_any_qkv( + pg: dist.ProcessGroup, + x: torch.Tensor, # (B, S_local, H, D) + *, + seq_lens: list[int], + use_sync: bool, +) -> tuple[torch.Tensor, int]: + """UAA forward all-to-all: (B, S_local, H, D) -> (B, S_global, H_local, D). + + Returns: + (resharded, orig_head_cnt) + """ + world_size = dist.get_world_size(pg) + if world_size == 1: + return x, int(x.shape[2]) + + bsz, s_local, head_cnt, head_dim = x.shape + orig_head_cnt = int(head_cnt) + padded_head_cnt = _ceil_div(orig_head_cnt, world_size) * world_size + head_pad = padded_head_cnt - orig_head_cnt + if head_pad: + x = F.pad(x, (0, 0, 0, head_pad)) + + head_cnt_local = padded_head_cnt // world_size + + # (B, S_local, H, D) -> (world_size, S_local, B, H_local, D) + x_t = x.reshape(bsz, s_local, world_size, head_cnt_local, head_dim).permute(2, 1, 0, 3, 4).contiguous() + # (world_size, S_local, B, H_local, D) -> (world_size * S_local, B, H_local, D) + x_t = x_t.flatten(0, 1) + + input_split_sizes = [s_local] * world_size + output_split_sizes = seq_lens + s_global = int(sum(output_split_sizes)) + + out = torch.empty((s_global, bsz, head_cnt_local, head_dim), device=x.device, dtype=x.dtype) + dist.all_to_all_single( + out, + x_t, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=pg, + ) + if use_sync: + from vllm_omni.platforms import current_omni_platform + + current_omni_platform.synchronize() + + # (S_global, B, H_local, D) -> (B, S_global, H_local, D) + out = out.permute(1, 0, 2, 3).contiguous() + return out, orig_head_cnt + + +def _ulysses_all_to_all_any_o( + pg: dist.ProcessGroup, + x: torch.Tensor, # (B, S_global, H_local, D) + *, + seq_lens: list[int], + local_seq_len: int, + orig_head_cnt: int, + use_sync: bool, +) -> torch.Tensor: + """UAA reverse all-to-all: (B, S_global, H_local, D) -> (B, S_local, H, D).""" + world_size = dist.get_world_size(pg) + if world_size == 1: + return x + + bsz, s_global, head_cnt_local, head_dim = x.shape + s_local = int(local_seq_len) + + # (B, S_global, H_local, D) -> (S_global, B, H_local, D) + x_t = x.permute(1, 0, 2, 3).contiguous() + + input_split_sizes = seq_lens + output_split_sizes = [s_local] * world_size + + out = torch.empty((world_size * s_local, bsz, head_cnt_local, head_dim), device=x.device, dtype=x.dtype) + dist.all_to_all_single( + out, + x_t, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=pg, + ) + if use_sync: + from vllm_omni.platforms import current_omni_platform + + current_omni_platform.synchronize() + + # (world_size * S_local, B, H_local, D) -> (B, S_local, H, D) + out = out.reshape(world_size, s_local, bsz, head_cnt_local, head_dim).permute(2, 1, 0, 3, 4).contiguous() + out = out.reshape(bsz, s_local, world_size * head_cnt_local, head_dim) + + if out.shape[2] != orig_head_cnt: + out = out[:, :, :orig_head_cnt, :].contiguous() + return out @dataclass(frozen=True, slots=True) @@ -24,6 +156,12 @@ class _UlyssesCtx(ParallelAttentionContext): use_sync: bool joint_len: int = 0 joint_strategy: str = "front" + # UAA (Ulysses Anything Attention) metadata + use_uaa: bool = False + uaa_seq_lens: tuple[int, ...] = () + uaa_local_seq_len: int = 0 + orig_head_cnt: int = 0 + joint_orig_head_cnt: int = 0 class UlyssesParallelAttention: @@ -65,9 +203,11 @@ class UlyssesParallelAttention: value: torch.Tensor, attn_metadata: AttentionMetadata | None, ): + mode = get_ulysses_mode(default="strict") joint_tensor_query = joint_tensor_key = joint_tensor_value = None joint_strategy = "front" joint_len = 0 + joint_orig_head_cnt = 0 if attn_metadata is not None: joint_tensor_query = attn_metadata.joint_query @@ -88,7 +228,27 @@ class UlyssesParallelAttention: # joint_query is (B, S, H, D). We split H (dim 2). ulysses_world_size = self._sp_group.ulysses_world_size ulysses_rank = self._sp_group.ulysses_rank - attn_heads_per_ulysses_rank = joint_tensor_query.shape[-2] // ulysses_world_size + joint_head_cnt = int(joint_tensor_query.shape[-2]) + joint_orig_head_cnt = joint_head_cnt + + if mode == "advanced_uaa": + padded_joint_head_cnt = _ceil_div(joint_head_cnt, ulysses_world_size) * ulysses_world_size + joint_head_pad = padded_joint_head_cnt - joint_head_cnt + if joint_head_pad: + joint_tensor_query = F.pad(joint_tensor_query, (0, 0, 0, joint_head_pad)) + joint_tensor_key = F.pad(joint_tensor_key, (0, 0, 0, joint_head_pad)) + joint_tensor_value = F.pad(joint_tensor_value, (0, 0, 0, joint_head_pad)) + joint_head_cnt = padded_joint_head_cnt + else: + if joint_head_cnt % ulysses_world_size != 0: + supported = _positive_divisors(joint_head_cnt) + raise ValueError( + "Ulysses-SP strict mode requires joint head_cnt divisible by ulysses_degree. " + f"joint_head_cnt={joint_head_cnt}, ulysses_degree={ulysses_world_size}. " + f"Try ulysses_degree in {supported}, or set ulysses_mode='advanced_uaa'." + ) + + attn_heads_per_ulysses_rank = joint_head_cnt // ulysses_world_size # Note: We use the same heads for Q/K/V joint_tensor_query = joint_tensor_query[ @@ -126,10 +286,55 @@ class UlyssesParallelAttention: attn_metadata.joint_key = joint_tensor_key attn_metadata.joint_value = joint_tensor_value - # (bs, seq_len/P, head_cnt, head_size) -> (bs, seq_len, head_cnt/P, head_size) - query = SeqAllToAll4D.apply(self._ulysses_pg, query, self._scatter_idx, self._gather_idx, self._use_sync) - key = SeqAllToAll4D.apply(self._ulysses_pg, key, self._scatter_idx, self._gather_idx, self._use_sync) - value = SeqAllToAll4D.apply(self._ulysses_pg, value, self._scatter_idx, self._gather_idx, self._use_sync) + ulysses_world_size = self._sp_group.ulysses_world_size + if mode == "advanced_uaa": + if self._scatter_idx != 2 or self._gather_idx != 1: + raise ValueError( + "ulysses_mode='advanced_uaa' currently only supports scatter_idx=2, gather_idx=1 " + f"(got scatter_idx={self._scatter_idx}, gather_idx={self._gather_idx})." + ) + + local_seq_len = int(query.shape[1]) + seq_lens = _all_gather_int(self._ulysses_pg, local_seq_len, device=query.device) + s_global = int(sum(seq_lens)) + + # In hybrid Ulysses+Ring, Ring attention uses P2P send/recv with fixed-shape + # buffers. This requires all ring ranks to have the same seq_len after the + # Ulysses all-to-all (i.e. per-ring-rank S_global must match). + if self._sp_group.ring_world_size > 1: + ring_s_globals = _all_gather_int(self._sp_group.ring_group, s_global, device=query.device) + if len(set(ring_s_globals)) != 1: + raise ValueError( + "ulysses_mode='advanced_uaa' with hybrid Ulysses+Ring requires the " + "post-Ulysses seq_len to be equal across ring ranks, but got " + f"{ring_s_globals} (ring_degree={self._sp_group.ring_world_size}). " + "This typically means the input sequence was not evenly shardable across the ring. " + "Try setting ring_degree=1, or choose a sequence length divisible by ring_degree." + ) + query, orig_head_cnt = _ulysses_all_to_all_any_qkv( + self._ulysses_pg, query, seq_lens=seq_lens, use_sync=self._use_sync + ) + key, _ = _ulysses_all_to_all_any_qkv(self._ulysses_pg, key, seq_lens=seq_lens, use_sync=self._use_sync) + value, _ = _ulysses_all_to_all_any_qkv(self._ulysses_pg, value, seq_lens=seq_lens, use_sync=self._use_sync) + else: + # Strict mode: fail fast with actionable errors for head divisibility. + for name, t in (("query", query), ("key", key), ("value", value)): + head_cnt = int(t.shape[2]) + if head_cnt % ulysses_world_size != 0: + supported = _positive_divisors(head_cnt) + raise ValueError( + "Ulysses-SP strict mode requires head_cnt divisible by ulysses_degree. " + f"{name}_head_cnt={head_cnt}, ulysses_degree={ulysses_world_size}. " + f"Try ulysses_degree in {supported}, or set ulysses_mode='advanced_uaa'." + ) + + # (bs, seq_len/P, head_cnt, head_size) -> (bs, seq_len, head_cnt/P, head_size) + query = SeqAllToAll4D.apply(self._ulysses_pg, query, self._scatter_idx, self._gather_idx, self._use_sync) + key = SeqAllToAll4D.apply(self._ulysses_pg, key, self._scatter_idx, self._gather_idx, self._use_sync) + value = SeqAllToAll4D.apply(self._ulysses_pg, value, self._scatter_idx, self._gather_idx, self._use_sync) + seq_lens = [] + local_seq_len = 0 + orig_head_cnt = 0 if is_joint: # Concatenate joint query AFTER AllToAll @@ -162,6 +367,11 @@ class UlyssesParallelAttention: use_sync=self._use_sync, joint_len=joint_len, joint_strategy=joint_strategy, + use_uaa=(mode == "advanced_uaa"), + uaa_seq_lens=tuple(int(x) for x in seq_lens) if mode == "advanced_uaa" else (), + uaa_local_seq_len=int(local_seq_len) if mode == "advanced_uaa" else 0, + orig_head_cnt=int(orig_head_cnt) if mode == "advanced_uaa" else 0, + joint_orig_head_cnt=int(joint_orig_head_cnt) if mode == "advanced_uaa" else 0, ) if attn_metadata is not None: @@ -217,7 +427,19 @@ class UlyssesParallelAttention: # SeqAllToAll4D handles: Scatter gather_idx, Gather scatter_idx. # Forward: Scatter 2 (H), Gather 1 (S). # Reverse: Scatter 1 (S), Gather 2 (H). - output_img = SeqAllToAll4D.apply(ctx.ulysses_pg, output_img, ctx.gather_idx, ctx.scatter_idx, ctx.use_sync) + if ctx.use_uaa: + output_img = _ulysses_all_to_all_any_o( + ctx.ulysses_pg, + output_img, + seq_lens=list(ctx.uaa_seq_lens), + local_seq_len=ctx.uaa_local_seq_len, + orig_head_cnt=ctx.orig_head_cnt, + use_sync=ctx.use_sync, + ) + else: + output_img = SeqAllToAll4D.apply( + ctx.ulysses_pg, output_img, ctx.gather_idx, ctx.scatter_idx, ctx.use_sync + ) # 2. Process Joint part: AllGather on Heads # Input: (B, JointLen, H/P, D). Output: (B, JointLen, H, D). @@ -227,6 +449,8 @@ class UlyssesParallelAttention: gathered_joint = [torch.zeros_like(output_joint) for _ in range(dist.get_world_size(ctx.ulysses_pg))] dist.all_gather(gathered_joint, output_joint, group=ctx.ulysses_pg) output_joint = torch.cat(gathered_joint, dim=2) + if ctx.use_uaa and ctx.joint_orig_head_cnt > 0 and output_joint.shape[2] != ctx.joint_orig_head_cnt: + output_joint = output_joint[:, :, : ctx.joint_orig_head_cnt, :].contiguous() # 3. Recombine if ctx.joint_strategy == "front": @@ -235,4 +459,13 @@ class UlyssesParallelAttention: return torch.cat([output_img, output_joint], dim=1) # Standard Ulysses Reverse + if ctx.use_uaa: + return _ulysses_all_to_all_any_o( + ctx.ulysses_pg, + attn_output, + seq_lens=list(ctx.uaa_seq_lens), + local_seq_len=ctx.uaa_local_seq_len, + orig_head_cnt=ctx.orig_head_cnt, + use_sync=ctx.use_sync, + ) return SeqAllToAll4D.apply(ctx.ulysses_pg, attn_output, ctx.gather_idx, ctx.scatter_idx, ctx.use_sync) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index a2b0d1d93..f50f2c0ea 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -55,6 +55,20 @@ class DiffusionParallelConfig: ring_degree: int = 1 """Number of GPUs used for ring sequence parallelism.""" + ulysses_mode: str = "strict" + """Ulysses sequence-parallel mode. + + - "strict": Require divisibility constraints (fastest, default). + - "advanced_uaa": Enable UAA ("Ulysses Anything Attention") to support + uneven sequence lengths and non-divisible head counts. + + Note: + - Ring attention does not support `attention_mask`, so models that rely on + mask-based auto-padding are still incompatible with Ring. + - When used in hybrid Ulysses+Ring, Ring requires consistent per-rank + sequence shapes across the ring group. + """ + cfg_parallel_size: int = 1 """Number of Classifier Free Guidance (CFG) parallel groups.""" @@ -88,6 +102,9 @@ class DiffusionParallelConfig: "Sequence parallel size must be equal to the product of ulysses degree and ring degree," f" but got {self.sequence_parallel_size} != {self.ulysses_degree} * {self.ring_degree}" ) + assert self.ulysses_mode in {"strict", "advanced_uaa"}, ( + f"ulysses_mode must be one of {{'strict','advanced_uaa'}}, but got {self.ulysses_mode!r}." + ) # Validate HSDP configuration if self.use_hsdp: diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index db574db1d..46b2b4a5c 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/vllm_omni/diffusion/forward_context.py @@ -80,6 +80,23 @@ def is_forward_context_available() -> bool: return _forward_context is not None +def get_ulysses_mode(*, default: str = "strict") -> str: + """Resolve the Ulysses-SP mode from the current ForwardContext. + + Returns `default` when ForwardContext is unavailable or the diffusion + config is not set. + """ + if not is_forward_context_available(): + return default + + cfg = get_forward_context().omni_diffusion_config + if cfg is None: + return default + + parallel_config = cfg.parallel_config + return str(getattr(parallel_config, "ulysses_mode", default)) + + def create_forward_context( vllm_config: VllmConfig | None = None, omni_diffusion_config: OmniDiffusionConfig | None = None, diff --git a/vllm_omni/diffusion/hooks/sequence_parallel.py b/vllm_omni/diffusion/hooks/sequence_parallel.py index 7044a3aa6..4233a29c0 100644 --- a/vllm_omni/diffusion/hooks/sequence_parallel.py +++ b/vllm_omni/diffusion/hooks/sequence_parallel.py @@ -315,10 +315,38 @@ class SequenceParallelSplitHook(ModelHook): logger.warning_once(f"Expected tensor with {sp_input.expected_dims} dims, got {x.dim()}. Skipping split.") return x + def _raise_strict_divisibility_error(*, dim: int, seq_len: int, sp_size: int) -> None: + # Keep message actionable: strict mode must be evenly shardable at the split hook level. + msg = ( + "Sequence Parallel strict mode requires the sequence length to be evenly shardable at split time, " + f"but got seq_len={seq_len} (dim={dim}) not divisible by sequence_parallel_size={sp_size}. " + ) + suggestions: list[str] = [] + if self.config.ulysses_degree > 1: + suggestions.append("set parallel_config.ulysses_mode='advanced_uaa' (UAA)") + suggestions.append( + "enable auto_pad=True in the model _sp_plan (requires attention_mask support and ring_degree=1)" + ) + suggestions.append("choose seq_len divisible by sequence_parallel_size (or adjust SP degrees)") + msg += "Try: " + "; ".join(suggestions) + "." + raise ValueError(msg) + + def _maybe_validate_strict_divisibility(*, dim: int, seq_len: int) -> None: + from vllm_omni.diffusion.forward_context import get_ulysses_mode + + if get_ulysses_mode(default="strict") != "strict": + return + sp_size = int(self.config.sequence_parallel_size) + if sp_size <= 1: + return + if seq_len < sp_size or (seq_len % sp_size) != 0: + _raise_strict_divisibility_error(dim=dim, seq_len=int(seq_len), sp_size=sp_size) + if isinstance(sp_input, SequenceParallelInput): # Full split with optional auto-padding if sp_input.auto_pad: return self._shard_with_auto_pad(x, sp_input.split_dim) + _maybe_validate_strict_divisibility(dim=sp_input.split_dim, seq_len=x.size(sp_input.split_dim)) return sp_shard(x, sp_input.split_dim, validate=False) elif isinstance(sp_input, SequenceParallelPartialInput): # Partial split: keep text portion, split image portion @@ -330,6 +358,7 @@ class SequenceParallelSplitHook(ModelHook): image_part = x.narrow(dim, text_len, x.size(dim) - text_len) # Only shard the image portion + _maybe_validate_strict_divisibility(dim=dim, seq_len=image_part.size(dim)) image_part_sharded = sp_shard(image_part, dim, validate=False) # Concatenate back: [text_full, image_sharded] diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index b7367e9b3..faeff3dce 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -623,6 +623,11 @@ class ZImageTransformer2DModel(CachedTransformer): quant_config: "QuantizationConfig | None" = None, ) -> None: super().__init__() + # NOTE: `DiffusersPipelineLoader.load_model()` initializes this module + # under `set_default_torch_dtype(od_config.dtype)`, so using the current + # default dtype keeps `self.dtype` consistent with the actually loaded + # weights. This dtype is used by the pipeline to cast inputs. + self.dtype = torch.get_default_dtype() self.in_channels = in_channels self.out_channels = in_channels self.all_patch_size = all_patch_size diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 75203bee7..f1ff9c1e4 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -802,6 +802,7 @@ class AsyncOmniEngine: if parallel_config is None: ulysses_degree = normalized_kwargs.get("ulysses_degree") or 1 ring_degree = normalized_kwargs.get("ring_degree") or 1 + ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict" sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size") tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1 cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1 @@ -819,6 +820,7 @@ class AsyncOmniEngine: sequence_parallel_size=sequence_parallel_size, ulysses_degree=ulysses_degree, ring_degree=ring_degree, + ulysses_mode=ulysses_mode, cfg_parallel_size=cfg_parallel_size, vae_patch_parallel_size=vae_patch_parallel_size, use_hsdp=use_hsdp, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 6830fc44b..f924d64c3 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -200,6 +200,15 @@ class OmniServeCommand(CLISubcommand): help="Ulysses Sequence Parallelism degree for diffusion models. " "Equivalent to setting DiffusionParallelConfig.ulysses_degree.", ) + omni_config_group.add_argument( + "--ulysses-mode", + type=str, + default="strict", + choices=["strict", "advanced_uaa"], + help="Ulysses sequence-parallel mode for diffusion models. " + "'strict' keeps the original divisibility requirements; " + "'advanced_uaa' enables the experimental UAA path for uneven sequence/head shapes.", + ) omni_config_group.add_argument( "--ring", "--ring-degree",