diff --git a/benchmarks/diffusion/bench_ulysses_cp_topology.py b/benchmarks/diffusion/bench_ulysses_cp_topology.py new file mode 100644 index 00000000000..a68528a04be --- /dev/null +++ b/benchmarks/diffusion/bench_ulysses_cp_topology.py @@ -0,0 +1,297 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Benchmark the dense compute-vs-communication tradeoff behind the composed +Ulysses x AllGather-KV topology, with a per-case profiler breakdown. + +Compares two 4-card sequence-parallel layouts for the SAME synthetic attention +shape, sweeping the sequence length: + + * ``usp4+uaa`` : pure Ulysses on 4 cards. With ``heads`` not divisible by 4, + advanced_uaa pads the head count (a real +14% attention + work when heads=14 -> padded 16), but communication is only + the Ulysses all-to-all. + * ``usp2+cp2`` : Ulysses(2) x AllGather-KV(2). No head padding (heads % 2 == 0 + for the default 14), but K/V are all-gathered over the + orthogonal group (extra communication). + +Why ``heads=14`` is the faithful 4-card analogue of MiniMax-H3's 56 heads on +16 cards (usp16+uaa vs usp8+cp2): both have heads/usp = 3.5, so the UAA pad is +56->64 (+14.3%) and 14->16 (+14.3%) respectively, and 56 % 8 == 0 mirrors +14 % 2 == 0 on the composed side. Running 56 heads on 4 cards would silently +disable the UAA padding (56 % 4 == 0) and stop measuring the effect of interest. + +For each (topology, seq) case the script reports three numbers from a single +``torch.profiler`` CUDA trace plus a wall-clock median: + + * ``total`` : median wall-clock ms of ``Attention.forward`` (iters runs). + * ``comm`` : summed CUDA device time of NCCL collectives (all-to-all / + all-gather), i.e. the communication cost. + * ``kernel`` : summed CUDA device time of everything else, i.e. the + attention kernel (flash/SDPA). + +Run on a box with >= 4 GPUs: + + python benchmarks/diffusion/bench_ulysses_cp_topology.py \ + --heads 14 --head-dim 128 --seqs 16384,32768,65536 --warmup 2 --iters 5 +""" + +from __future__ import annotations + +import argparse +import json +import os +import tempfile +import time + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch.profiler import ProfilerActivity, profile + +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.config import set_current_diffusion_config +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 set_forward_context +from vllm_omni.platforms import current_omni_platform + +# One entry per topology: (label, ulysses_degree, allgather_degree, ulysses_mode). +_TOPOLOGIES = [ + ("usp4+uaa", 4, 1, "advanced_uaa"), + ("usp2+cp2", 2, 2, "strict"), +] + +# NCCL collective kernels are prefixed with "nccl" in the CUDA trace; that is +# what we count as communication. Everything else is the attention kernel. +_COMM_KEYWORDS = ("nccl", "all_to_all", "all_gather") + + +def _profile_once(attn: Attention, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> tuple[float, float, list]: + """Run one forward under a CUDA profiler. + + Returns ``(comm_ms, kernel_ms, comm_ops)`` where ``comm_ops`` lists every + NCCL collective with its summed device time. Every rank must call this + (the forward contains collectives); only rank 0 uses the result. + """ + with profile(activities=[ProfilerActivity.CUDA]) as prof: + attn(q, k, v) + current_omni_platform.synchronize() + + comm_us = 0.0 + kernel_us = 0.0 + comm_ops: dict[str, float] = {} + for evt in prof.key_averages(): + # torch >= 2.x names this device_time_total; older builds used cuda_time_total. + t = getattr(evt, "device_time_total", 0.0) or getattr(evt, "cuda_time_total", 0.0) + if t <= 0: + continue + name = evt.key + if any(k in name.lower() for k in _COMM_KEYWORDS): + comm_us += t + comm_ops[name] = comm_ops.get(name, 0.0) + t / 1000.0 + else: + kernel_us += t + return comm_us / 1000.0, kernel_us / 1000.0, sorted(comm_ops.items(), key=lambda x: -x[1]) + + +def _worker( + rank: int, + world_size: int, + topology: tuple, + heads: int, + head_dim: int, + seqs: list[int], + dtype: torch.dtype, + warmup: int, + iters: int, + result_file: str, +) -> None: + label, ulysses_degree, allgather_degree, ulysses_mode = topology + + os.environ.update( + { + "RANK": str(rank), + "LOCAL_RANK": str(rank), + "WORLD_SIZE": str(world_size), + "MASTER_ADDR": "localhost", + "MASTER_PORT": "12345", + } + ) + device = torch.device(f"{current_omni_platform.device_type}:{rank}") + current_omni_platform.set_device(device) + torch.set_default_dtype(dtype) + + init_distributed_environment() + + 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=1, + allgather_degree=allgather_degree, + ulysses_mode=ulysses_mode, + cfg_parallel_size=1, + ) + od_config = OmniDiffusionConfig.from_kwargs( + model="benchmark", + dtype=dtype, + parallel_config=parallel_config, + # Force one kernel for both topologies so only the parallel path differs. + diffusion_attention_backend="TORCH_SDPA", + ) + initialize_model_parallel( + data_parallel_size=1, + cfg_parallel_size=1, + sequence_parallel_size=world_size, + ulysses_degree=ulysses_degree, + ring_degree=1, + allgather_degree=allgather_degree, + tensor_parallel_size=1, + pipeline_parallel_size=1, + ) + + # Build the attention layer and run the whole sweep inside the forward + # context: get_ulysses_mode() reads ulysses_mode from it at forward time, + # so leaving the context would silently drop advanced_uaa back to strict. + with set_forward_context(omni_diffusion_config=od_config), set_current_diffusion_config(od_config): + attn = Attention( + num_heads=heads, + head_size=head_dim, + causal=False, + softmax_scale=1.0 / (head_dim**0.5), + num_kv_heads=None, # MHA, matching MiniMax-H3 (no num_key_value_heads) + scatter_idx=2, + gather_idx=1, + use_sync=False, + ) + + medians: dict[int, float] = {} + profiles: dict[int, dict] = {} + for seq in seqs: + assert seq % world_size == 0, f"seq {seq} must be divisible by world_size {world_size}" + s_local = seq // world_size + q = torch.randn(1, s_local, heads, head_dim, device=device, dtype=dtype) + k = torch.randn(1, s_local, heads, head_dim, device=device, dtype=dtype) + v = torch.randn(1, s_local, heads, head_dim, device=device, dtype=dtype) + + def _run() -> None: + # NOTE: every rank must call forward (collectives); only rank 0 clocks. + attn(q, k, v) + + for _ in range(warmup): + _run() + current_omni_platform.synchronize() + + times = [] + for _ in range(iters): + current_omni_platform.synchronize() + t0 = time.perf_counter() + _run() + current_omni_platform.synchronize() + times.append((time.perf_counter() - t0) * 1000.0) + times.sort() + medians[seq] = times[len(times) // 2] + + # One profiled forward per rank. A straggler rank inflates every + # other rank's measured collective time, so take the rank with the + # smallest comm time as the true communication cost. + comm_ms, kernel_ms, comm_ops = _profile_once(attn, q, k, v) + all_splits = [None] * world_size + dist.all_gather_object(all_splits, (comm_ms, kernel_ms, comm_ops)) + if rank == 0: + best_comm, best_kernel, best_ops = min(all_splits, key=lambda x: x[0]) + profiles[seq] = {"comm": best_comm, "kernel": best_kernel, "comm_ops": best_ops} + + destroy_distributed_env() + if rank == 0: + with open(result_file, "w", encoding="utf-8") as f: + json.dump( + { + "label": label, + # JSON object keys are strings; restored to int on load. + "medians": {str(seq): ms for seq, ms in medians.items()}, + "profiles": {str(seq): prof for seq, prof in profiles.items()}, + }, + f, + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--heads", type=int, default=14) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--seqs", type=str, default="1024,2048,4096,8192,16384") + parser.add_argument("--dtype", type=str, default="bfloat16") + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--iters", type=int, default=10) + args = parser.parse_args() + + seqs = [int(x) for x in args.seqs.split(",")] + dtype = getattr(torch, args.dtype) + world_size = 4 + device_count = current_omni_platform.get_device_count() + if device_count < world_size: + raise SystemExit(f"Need {world_size} GPUs, found {device_count}.") + + # Sanity: if heads % 4 == 0, usp4+uaa degrades to strict (no padding) and + # the comparison stops measuring the effect of interest. + if args.heads % 4 == 0: + print( + f"[warn] heads={args.heads} is divisible by 4, so usp4+uaa will NOT pad " + f"(use a heads value with heads % 4 != 0 and heads % 2 == 0, e.g. 14, " + f"to reproduce the 56-head / usp16+uaa tradeoff faithfully)." + ) + + results: dict[str, tuple[dict[int, float], dict[int, dict]]] = {} + for topology in _TOPOLOGIES: + label = topology[0] + with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: + result_file = f.name + try: + mp.spawn( + _worker, + args=( + world_size, + topology, + args.heads, + args.head_dim, + seqs, + dtype, + args.warmup, + args.iters, + result_file, + ), + nprocs=world_size, + join=True, + ) + with open(result_file, encoding="utf-8") as f: + payload = json.load(f) + got_label = payload["label"] + assert got_label == label, (got_label, label) + medians = {int(seq): ms for seq, ms in payload["medians"].items()} + profiles = {int(seq): prof for seq, prof in payload["profiles"].items()} + results[label] = (medians, profiles) + finally: + if os.path.exists(result_file): + os.remove(result_file) + + print(f"\nheads={args.heads} head_dim={args.head_dim} dtype={args.dtype} warmup={args.warmup} iters={args.iters}") + for seq in seqs: + print(f"\n=== seq={seq} ===") + for label, *_ in _TOPOLOGIES: + medians, profiles = results[label] + total = medians[seq] + p = profiles[seq] + print(f" {label:<10} total={total:>8.3f}ms comm={p['comm']:>8.3f}ms kernel={p['kernel']:>8.3f}ms") + for op_name, op_ms in p["comm_ops"]: + print(f" {op_name:<50} {op_ms:>8.3f}ms") + + +if __name__ == "__main__": + main() diff --git a/docs/user_guide/diffusion/parallelism/sequence_parallel.md b/docs/user_guide/diffusion/parallelism/sequence_parallel.md index 9bb5b5d7daf..47af53a560e 100644 --- a/docs/user_guide/diffusion/parallelism/sequence_parallel.md +++ b/docs/user_guide/diffusion/parallelism/sequence_parallel.md @@ -1,6 +1,5 @@ # Sequence Parallelism Guide - ## Table of Content - [Overview](#overview) @@ -23,6 +22,7 @@ See supported models list in [Diffusion Features - Supported Models](../../diffu - **DeepSpeed Ulysses Sequence Parallel (Ulysses-SP)** ([paper](https://arxiv.org/pdf/2309.14509)): Uses all-to-all communication for subset of attention heads per device - **Ring-Attention** ([paper](https://arxiv.org/abs/2310.01889)): Uses ring-based P2P communication with sharded sequence dimension throughout - **Hybrid Ulysses + Ring**: Combines both for larger scale parallelism (`ulysses_degree × ring_degree`) +- **Ulysses × AllGather-KV**: A two-dimensional topology where Ulysses shards attention heads and AllGather-KV makes K/V global (`ulysses_degree × allgather_degree`) --- @@ -85,6 +85,54 @@ omni = Omni( ) ``` +**Ulysses × AllGather-KV** (4 GPUs total, no Ring in the attention computation): + +```python +omni = Omni( + model="Qwen/Qwen-Image", + parallel_config=DiffusionParallelConfig(ulysses_degree=2, allgather_degree=2) +) +``` + +### Ulysses × AllGather-KV + +This topology reaches larger SP degrees without putting Ring Attention on the +critical path. Each rank starts from `[B, S / (U × A), H, D]` and the forward is: + +```text +Ulysses all-to-all (U group): Q, K, V → [B, S / A, H / U, D] +K/V AllGather (A group): K, V → [B, S, H / U, D] +local-Q / global-KV attention: O → [B, S / A, H / U, D] +reverse Ulysses all-to-all: O → [B, S / (U × A), H, D] +``` + +The kernel always sees a local-Q/global-KV problem, so the full global sequence +is materialized in a single, ordered K/V tensor. That matters for attention +backends that need stable global K/V indexing — sparse-mask generation in +particular — which is awkward under Ring Attention because Ring keeps K/V +block-local and owns distributed attention itself. + +Choosing between the two 4-GPU options is mostly about the attention backend: if +the backend can consume a global K/V tensor, prefer Ulysses × AllGather-KV; if +the sequence is large enough that materializing global K/V does not fit, use +Ulysses × Ring instead. + +Constraints: + +- `ring_degree` must be `1` — composing all three SP dimensions is not supported. +- Non-causal attention only. +- Every rank's region must be equally long, so the shared sequence must be evenly + shardable across the SP group (or the model's `_sp_plan` must use `auto_pad`). + A mismatch fails fast instead of corrupting the collective. +- 2D key masks are supported: an image mask must cover the gathered *global* + image keys replicated across SP ranks (the same contract as plain + AllGather-KV), and a joint text-padding mask (`joint_attn_mask`, e.g. + Qwen-Image with unequal prompt lengths) is merged automatically in + `joint_strategy` order. Masks covering only the local shard fail fast. + `mask_sp_padding=True` builds exactly such a global 2D mask, so strict + padding masking works under this topology; 4D masks keep query-range slicing. +- Not supported with Scheduler-paged KV or with `diffusion_attention_backend="TRTLLM_ATTN"`. + --- ## Example Script @@ -124,6 +172,17 @@ python examples/offline_inference/text_to_image/text_to_image.py \ --width 1024 --height 1024 ``` +**Ulysses × AllGather-KV:** + +```bash +# 2 Ulysses × 2 AllGather-KV = 4 GPUs total, ring_degree stays 1 +python examples/offline_inference/text_to_image/text_to_image.py \ + --model Qwen/Qwen-Image \ + --prompt "A cat sitting on a windowsill" \ + --ulysses-degree 2 --allgather-degree 2 \ + --width 1024 --height 1024 +``` + ### Online Serving **Ulysses-SP:** @@ -153,6 +212,13 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2 vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2 ``` +**Ulysses × AllGather-KV:** + +```bash +# Text-to-image (requires >= 4 GPUs); --ring must stay 1 +vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --allgather-degree 2 +``` + --- ## Configuration Parameters @@ -160,18 +226,19 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2 In `DiffusionParallelConfig`: | Parameter | Type | Default | Description | -|-----------|------|---------|-------------| +| ----------- | ------ | --------- | ------------- | | `ulysses_degree` | int | 1 | Number of GPUs for Ulysses-SP. Uses all-to-all communication. | | `ring_degree` | int | 1 | Number of GPUs for Ring-Attention. Uses P2P ring communication. | +| `allgather_degree` | int | 1 | Number of GPUs for AllGather-KV sequence parallelism. Gathers K/V in an orthogonal group so every rank runs local-Q/global-KV attention. May be combined with `ulysses_degree`, but not with `ring_degree > 1`. | | `ulysses_mode` | str | `"default"` | Ulysses attention mode. Set to `"advanced_uaa"` to handle arbitrary sequence lengths and head counts without padding. | | `mask_sp_padding` | bool | `False` | When the sequence length is not divisible by the SP size, tokens are auto-padded with zeros. Set to `True` to mask those padding tokens (strict, but uses the slower varlen attention path); the default `False` leaves them unmasked, keeping the fast path with negligible numerical impact. | **Notes:** -- Total sequence parallel size equals to `ulysses_degree × ring_degree` + +- Total sequence parallel size equals `ulysses_degree × ring_degree × allgather_degree` - Degrees must evenly divide the sequence length for optimal performance (or use `ulysses_mode="advanced_uaa"` for Ulysses-SP) - `mask_sp_padding` is an experimental feature, currently only supported by `Wan2.2`, `Wan2.2 Vace`, `Qwen-Image`, `Flux 2`, and `HunyuanVideo 1.5` - ## Best Practices ### When to Use @@ -185,7 +252,6 @@ In `DiffusionParallelConfig`: - Small images (<1024px) - overhead exceeds benefit, use single GPU with cache instead - --- ## Troubleshooting @@ -195,6 +261,7 @@ In `DiffusionParallelConfig`: **Symptoms**: Adding GPUs doesn't improve speed proportionally, or higher parallelism degree is slower **Diagnosis:** + ```bash # Check GPU topology nvidia-smi topo -m @@ -205,16 +272,17 @@ nvidia-smi topo -m 1. Check inter-GPU communication - NVLink is better than PCIe 2. Reduce parallelism degree if over-parallelized: + ```python # If 4 GPUs is slower than 2 parallel_config=DiffusionParallelConfig(ulysses_degree=2) ``` -3. Try to switch between Ring-Attention and Ulysses-SP + +1. Try to switch between Ring-Attention and Ulysses-SP - Ring-Attention has advantages, like communication-computation overlap, but the block-wise loop overhead is relatively higher, especially for short sequences - Ulysses-SP: can benefit from larger bandwidth (such as NVLink), with two major constraints, the sequence length should be divisible by usp size, and the number of heads should be divisible by usp size (or use `ulysses_mode="advanced_uaa"`) - ### Common Issue 2: Out of Memory (OOM) **Symptoms**: CUDA OOM errors or process crashes with memory errors @@ -222,11 +290,12 @@ parallel_config=DiffusionParallelConfig(ulysses_degree=2) **Solutions:** 1. Increase parallelism degree to split sequence more: + ```python parallel_config=DiffusionParallelConfig(ulysses_degree=4) # From 2 ``` -2. Combine with other parallelism method, e.g., tensor parallel, and memory optimization methods, e.g., cpu offloading. +1. Combine with other parallelism method, e.g., tensor parallel, and memory optimization methods, e.g., cpu offloading. ## Summary diff --git a/docs/user_guide/feature_compatibility.md b/docs/user_guide/feature_compatibility.md index eeb087510ad..4a6416bf75f 100644 --- a/docs/user_guide/feature_compatibility.md +++ b/docs/user_guide/feature_compatibility.md @@ -154,7 +154,6 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \ --ring 2 ``` - ## Limitations ### Incompatibilities @@ -174,10 +173,10 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \ ### Configuration Constraints - **GPU Count Must Match Parallel Degrees**: Total GPU count must satisfy: + ```text + total_gpus = ulysses_degree × ring_degree × allgather_degree × cfg_parallel_size × tensor_parallel_size ``` - total_gpus = ulysses_degree × ring_degree × cfg_parallel_size × tensor_parallel_size - ``` - Any mismatch will cause a configuration error at startup. + Any mismatch will cause a configuration error at startup. `allgather_degree > 1` may be combined with `ulysses_degree > 1`, but not with `ring_degree > 1`. - **VAE Patch Parallel Size ≤ DiT Process Group Size**: `vae_patch_parallel_size` reuses the DiT process group and cannot exceed it. Larger values are automatically clamped with a warning. @@ -190,6 +189,7 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \ **Symptoms:** Adding more GPUs doesn't improve speed proportionally **Solutions:** + 1. Check GPU communication bandwidth (use `nvidia-smi topo -m`) 2. Reduce parallelism degree if communication overhead is high 3. For very long sequences, prefer Ring-Attention over Ulysses-SP @@ -200,6 +200,7 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \ **Symptoms:** OOM errors when combining methods **Solutions:** + 1. Enable Tensor Parallelism to shard weights 2. Reduce resolution or batch size 3. Combine with memory efficient methods, such as cpu offloading @@ -209,6 +210,7 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \ **Symptoms:** Errors about invalid parallel configuration **Solutions:** + 1. Verify total GPU count matches: `ulysses × ring × cfg × tp` 2. Check model supports all enabled methods 3. Ensure divisibility constraints (e.g., Z-Image TP=1 or 2 only) diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 26b85a39cd8..31c190ccef6 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -1364,9 +1364,24 @@ def test_diffusion_parallel_config_accepts_four_way_guidance_parallelism(): assert cfg.world_size == 4 -def test_diffusion_parallel_config_rejects_allgather_with_ulysses_or_ring(): - with pytest.raises(ValidationError): - OmniStageDiffusionParallelConfig(allgather_degree=2, ulysses_degree=2) +def test_diffusion_parallel_config_accepts_ulysses_allgather_composition(): + cfg = OmniStageDiffusionParallelConfig(allgather_degree=2, ulysses_degree=2) + + assert cfg.sequence_parallel_size == 4 + assert cfg.world_size == 4 + + +def test_diffusion_parallel_config_rejects_allgather_composed_with_ring(): + with pytest.raises(ValidationError, match="cannot be composed with ring_degree"): + OmniStageDiffusionParallelConfig(allgather_degree=2, ring_degree=2) + + +def test_diffusion_parallel_config_derives_sequence_parallel_size_from_all_three_degrees(): + cfg = OmniStageDiffusionParallelConfig(ulysses_degree=2, ring_degree=2, allgather_degree=1) + assert cfg.sequence_parallel_size == 4 + + cfg = OmniStageDiffusionParallelConfig(ulysses_degree=2, allgather_degree=4) + assert cfg.sequence_parallel_size == 8 def test_stage_realizations_use_stage_specific_parallel_config_types(): diff --git a/tests/diffusion/attention/test_attention_sp.py b/tests/diffusion/attention/test_attention_sp.py index 727bb55cbe7..730d2fa40e0 100644 --- a/tests/diffusion/attention/test_attention_sp.py +++ b/tests/diffusion/attention/test_attention_sp.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for Ulysses + Ring sequence-parallel attention correctness. @@ -7,14 +7,24 @@ -------------- * ``test_sequence_parallel`` verifies that the ``Attention`` layer produces numerically equivalent results whether the sequence is processed on a single - rank (baseline, SP=1) or sharded across multiple ranks via Ulysses/Ring SP. - The test spawns two separate multi-process runs with ``torch.multiprocessing.spawn``: - 1. Baseline – world_size=1, ulysses_degree=1, ring_degree=1. - 2. SP run – world_size=ulysses_degree*ring_degree, each rank holds a - contiguous slice of the full sequence. + rank (baseline, SP=1) or sharded across multiple ranks via Ulysses / Ring / + AllGather-KV SP. The test spawns two separate multi-process runs with + ``torch.multiprocessing.spawn``: + 1. Baseline – world_size=1, no SP. + 2. SP run – world_size=ulysses_degree*ring_degree*allgather_degree, each + rank holds a contiguous slice of the full sequence. After both runs, rank-0 output tensors are compared element-wise with a tolerance appropriate for the dtype (bfloat16). +* The ``_Mock*`` tests pin the communication contract of the parallel attention + strategies on CPU, without a GPU or a process group. For the composed + Ulysses x AllGather-KV strategy in particular they assert the *ordering* of + the collectives (Ulysses first, then the K/V AllGather, then joint re-attach), + which the multi-process equivalence test cannot observe on its own: softmax + attention is invariant to a permutation of the K/V slots, so a wrong region + order only shows up when a mask or sparse index makes the slot position + meaningful. + SP-plan hooks are NOT applied in this test ------------------------------------------ ``ForwardContext.sp_plan_hooks_applied`` remains ``False``. As a result, @@ -27,16 +37,22 @@ import os import pickle import tempfile +from types import SimpleNamespace import pytest import torch +import torch.distributed from tests.helpers.mark import hardware_test -from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, QueryRange +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, PackedPaddingMetadata, QueryRange from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.attention.parallel.allgather_kv import ( AllGatherKVParallelAttention, ) +from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention +from vllm_omni.diffusion.attention.parallel.ulysses_allgather import ( + UlyssesAllGatherKVParallelAttention, +) from vllm_omni.diffusion.config import set_current_diffusion_config from vllm_omni.diffusion.data import ( DiffusionParallelConfig, @@ -64,6 +80,14 @@ def seed_everything(seed: int): current_omni_platform.manual_seed(seed) +def _sp_kind_label(ulysses_degree: int, ring_degree: int, allgather_degree: int) -> str: + if allgather_degree > 1 and ulysses_degree > 1: + return f"ulysses={ulysses_degree}, allgather={allgather_degree}" + if allgather_degree > 1: + return f"allgather={allgather_degree}" + return f"ulysses={ulysses_degree}, ring={ring_degree}" + + class TestAttentionModel(torch.nn.Module): """Test model using Attention layer.""" @@ -231,6 +255,149 @@ def test_allgather_kv_slices_full_dense_mask_to_local_query_rows(): assert torch.equal(metadata_local.attn_mask, expected_rows) +def _packed_metadata(cu: list[int], max_len: int, valid_kv_length: int | None = None, **kwargs): + extra = { + "cu_seqlens_q": torch.tensor(cu, dtype=torch.int32), + "cu_seqlens_k": torch.tensor(cu, dtype=torch.int32), + "max_seqlen_q": max_len, + "max_seqlen_k": max_len, + } + if valid_kv_length is not None: + extra["valid_kv_length"] = valid_kv_length + return AttentionMetadata(extra=extra, **kwargs) + + +def test_allgather_kv_remaps_packed_query_boundaries(): + """Packed backends consume cu_seqlens_q, not query_ranges: remap per rank. + + The review's example: two packed requests of 8 and 4 tokens, global + boundaries [0, 8, 12], allgather_degree=2. Rank 0 holds Q rows 0-5 of + request A; rank 1 holds A6 A7 | B0..B3. Each document keeps one local + entry (zero-length on rank 0 for B) so the one-to-one pairing with the + global cu_seqlens_k documents survives. + """ + img_seq_local = 6 + img_seq_full = 12 + cu = [0, 8, 12] + expected = {0: [0, 6, 6], 1: [0, 2, 6]} + expected_max_q = {0: 6, 1: 4} # exact per-rank maxima, for reference + + for rank in (0, 1): + key_chunks = [torch.zeros((1, img_seq_local, 1, 1)) for _ in range(2)] + value_chunks = [torch.zeros_like(chunk) for chunk in key_chunks] + strategy = AllGatherKVParallelAttention( + _MockAllGatherSPGroup(rank=rank, gather_chunks=[key_chunks, value_chunks]), + ) + metadata = _packed_metadata(cu, max_len=8) + query = torch.zeros((1, img_seq_local, 1, 1)) + + q_local, k_full, _, metadata_local, _ = strategy.pre_attention( + query, key_chunks[rank], value_chunks[rank], metadata + ) + + assert q_local.shape[1] == img_seq_local + assert k_full.shape[1] == img_seq_full + assert metadata_local is not None + assert metadata_local.query_ranges is not None + assert metadata_local.extra["cu_seqlens_q"].tolist() == expected[rank] + # Key side stays global: K/V were gathered. + assert metadata_local.extra["cu_seqlens_k"].tolist() == cu + # max_seqlen_q keeps the producer bound (a safe upper bound of every + # local segment, e.g. >= expected_max_q[rank], without a host sync). + assert metadata_local.extra["max_seqlen_q"] >= expected_max_q[rank] + assert metadata_local.extra["max_seqlen_q"] == 8 + + +def test_allgather_kv_packed_document_without_local_rows_keeps_entry(): + """A document fully outside this rank's span keeps a zero-length entry. + + Docs [0, 4, 12] with allgather_degree=2: rank 1's span [6, 12) holds no + row of document 0, which must still appear as a zero-length entry so the + Q/K document pairing stays one-to-one. + """ + img_seq_local = 6 + rank = 1 + cu = [0, 4, 12] + key_chunks = [torch.zeros((1, img_seq_local, 1, 1)) for _ in range(2)] + value_chunks = [torch.zeros_like(chunk) for chunk in key_chunks] + strategy = AllGatherKVParallelAttention( + _MockAllGatherSPGroup(rank=rank, gather_chunks=[key_chunks, value_chunks]), + ) + metadata = _packed_metadata(cu, max_len=8) + + _, _, _, metadata_local, _ = strategy.pre_attention( + torch.zeros((1, img_seq_local, 1, 1)), key_chunks[rank], value_chunks[rank], metadata + ) + + assert metadata_local is not None + local_cu = metadata_local.extra["cu_seqlens_q"].tolist() + assert local_cu == [0, 0, 6], "document 0 must keep a zero-length local entry" + assert metadata_local.extra["cu_seqlens_k"].tolist() == cu + + +def test_allgather_kv_packed_single_request_with_padding(): + """The [real, pad] single-document packing narrows to the local valid prefix. + + Global packing [0, 10] with 8 valid rows (valid_kv_length=8) and + allgather_degree=2: each rank owns 5 rows, of which rank 0 has 5 valid + and rank 1 has 3. packed_padding narrows q_length accordingly while the + key side keeps the global valid prefix. + """ + img_seq_local = 5 + cu = [0, 10] + for rank, local_valid in ((0, 5), (1, 3)): + key_chunks = [torch.zeros((1, img_seq_local, 1, 1)) for _ in range(2)] + value_chunks = [torch.zeros_like(chunk) for chunk in key_chunks] + strategy = AllGatherKVParallelAttention( + _MockAllGatherSPGroup(rank=rank, gather_chunks=[key_chunks, value_chunks]), + ) + cu_q = torch.tensor(cu, dtype=torch.int32) + metadata = _packed_metadata( + cu, + max_len=8, + valid_kv_length=8, + packed_padding=PackedPaddingMetadata( + q_length=8, + kv_length=8, + cu_seqlens_q=cu_q[:2].clone(), + cu_seqlens_k=cu_q[:2].clone(), + ), + ) + + _, _, _, metadata_local, _ = strategy.pre_attention( + torch.zeros((1, img_seq_local, 1, 1)), key_chunks[rank], value_chunks[rank], metadata + ) + + assert metadata_local is not None + assert metadata_local.extra["cu_seqlens_q"].tolist() == [0, img_seq_local] + packed = metadata_local.packed_padding + assert packed is not None + assert packed.q_length == local_valid + assert packed.kv_length == 8, "key side keeps the global valid prefix" + assert packed.cu_seqlens_q.tolist() == [0, img_seq_local] + assert packed.cu_seqlens_k.tolist() == [0, 10] + + +def test_allgather_kv_rejects_packed_metadata_with_joint(): + """Packed cu_seqlens combined with joint attention is undefined; fail closed.""" + rank = 1 + joint_len = 1 + img_seq_local = 2 + key_chunks = [torch.zeros((1, img_seq_local, 1, 1)) for _ in range(2)] + value_chunks = [torch.zeros_like(chunk) for chunk in key_chunks] + strategy = AllGatherKVParallelAttention( + _MockAllGatherSPGroup(rank=rank, gather_chunks=[key_chunks, value_chunks]), + ) + joint = torch.ones((1, joint_len, 1, 1)) + metadata = _packed_metadata([0, 4], max_len=4) + metadata.joint_query = joint + metadata.joint_key = joint + metadata.joint_value = joint + + with pytest.raises(NotImplementedError, match="packed cu_seqlens_q combined with joint"): + strategy.pre_attention(torch.zeros((1, img_seq_local, 1, 1)), key_chunks[rank], value_chunks[rank], metadata) + + def test_allgather_kv_slices_rear_joint_dense_mask(): rank = 1 joint_len = 1 @@ -410,6 +577,393 @@ def test_allgather_kv_keeps_gathered_kv_compressed_for_gqa(): torch.testing.assert_close(v_full, expected_value) +class _MockComposedSPGroup(_MockAllGatherSPGroup): + """AllGather SP-group stub extended with the Ulysses sub-group fields.""" + + def __init__(self, *, rank: int, allgather_world_size: int, gather_chunks: list[list[torch.Tensor]]) -> None: + super().__init__(rank=rank, gather_chunks=gather_chunks) + self.allgather_world_size = allgather_world_size + self.ulysses_group = object() + self.ulysses_world_size = 2 + self.ulysses_rank = rank % 2 + + +def test_composed_strategy_orders_ulysses_before_allgather(monkeypatch): + """K/V must be gathered *after* the Ulysses reshard, and joint K/V after that. + + Gathering pre-Ulysses shards would concatenate the strided head layout of + different regions (a permuted sequence), and folding joint K/V in before the + gather would replicate them once per AllGather rank. Both are pinned here by + the recorded gather shapes and the final key length. + """ + rank = 1 + img_local = 2 # S / (U * A), the rank's own shard + region = 4 # S / A, after the Ulysses all-to-all + heads_local = 2 + joint_len = 1 + head_dim = 1 + + key_chunks = [torch.full((1, region, heads_local, head_dim), 10.0 + i) for i in range(2)] + value_chunks = [torch.full((1, region, heads_local, head_dim), 20.0 + i) for i in range(2)] + sp_group = _MockComposedSPGroup(rank=rank, allgather_world_size=2, gather_chunks=[key_chunks, value_chunks]) + + seen: dict[str, object] = {} + + def fake_pre_attention(self, query, key, value, attn_metadata, defer_joint=False): + seen["defer_joint"] = defer_joint + seen["ulysses_input_seq_len"] = query.shape[1] + joint = torch.ones(1, joint_len, heads_local, head_dim) + if attn_metadata is not None: + attn_metadata.joint_query = joint + attn_metadata.joint_key = joint + attn_metadata.joint_value = joint + return ( + torch.zeros(1, region, heads_local, head_dim), + torch.full((1, region, heads_local, head_dim), 10.0 + rank), + torch.full((1, region, heads_local, head_dim), 20.0 + rank), + attn_metadata, + "ulysses-ctx", + ) + + def fake_post_attention(self, attn_output, ctx): + seen["post_ctx"] = ctx + return attn_output + + monkeypatch.setattr(UlyssesParallelAttention, "pre_attention", fake_pre_attention) + monkeypatch.setattr(UlyssesParallelAttention, "post_attention", fake_post_attention) + + strategy = UlyssesAllGatherKVParallelAttention(sp_group, scatter_idx=2, gather_idx=1, use_sync=False) + # Skip the one-time region-length collective: it is covered separately. + strategy._validated_region_len = region + + query = torch.zeros(1, img_local, 4, head_dim) + key = torch.zeros(1, img_local, 4, head_dim) + value = torch.zeros(1, img_local, 4, head_dim) + + q_out, k_out, v_out, _, ctx = strategy.pre_attention(query, key, value, AttentionMetadata(joint_strategy="front")) + + assert seen["defer_joint"] is True + assert seen["ulysses_input_seq_len"] == img_local, "Ulysses must see the local shard, not the gathered one" + # The gather consumed post-Ulysses K/V, exactly once each, on the AllGather group. + assert sp_group.gathered_input_shapes == [(1, region, heads_local, head_dim)] * 2 + # Joint K/V are re-attached after the gather, so they appear exactly once. + assert k_out.shape[1] == joint_len + 2 * region + assert v_out.shape[1] == joint_len + 2 * region + assert q_out.shape[1] == joint_len + region + torch.testing.assert_close(k_out[:, joint_len:], torch.cat(key_chunks, dim=1)) + torch.testing.assert_close(v_out[:, joint_len:], torch.cat(value_chunks, dim=1)) + + # The reverse transform stays Ulysses', on the same ctx. + assert strategy.post_attention(q_out, ctx) is q_out + assert seen["post_ctx"] == "ulysses-ctx" + assert strategy.name == "ulysses_allgather_kv" + + +def test_composed_strategy_normalizes_image_only_mask(monkeypatch): + """A 2D image-only mask must cover the gathered *global* image keys. + + With ``_merge_joint_attn_mask`` owning all 2D key-mask handling there is + no blanket 2D rejection anymore: the mask is validated against the + rebuilt global image layout and normalized in place. A shard-local or + stale-length mask fails closed. + """ + region = 4 + make = _composed_mask_mocks(monkeypatch, region=region, joint_len=1, joint_fields=()) + query = torch.zeros(1, 2, 4, 1) + + img_mask = torch.tensor([[True, True, False, True, True, True, True, True]]) + _, strategy = make() + _, k_out, _, merged, _ = strategy.pre_attention(query, query, query, AttentionMetadata(attn_mask=img_mask)) + + assert k_out.shape[1] == 2 * region + assert merged.attn_mask is not None + assert merged.attn_mask.shape == (1, 2 * region) + assert merged.attn_mask.dtype == torch.bool + torch.testing.assert_close(merged.attn_mask, img_mask) + + _, strategy = make() + with pytest.raises(ValueError, match="global image keys"): + strategy.pre_attention( + query, query, query, AttentionMetadata(attn_mask=torch.ones(1, 2 * region - 1, dtype=torch.bool)) + ) + + +def _composed_mask_mocks( + monkeypatch, + *, + rank: int = 1, + region: int = 4, + heads_local: int = 2, + joint_len: int = 1, + joint_fields: tuple[str, ...] = ("query", "key", "value"), +): + """Wire the composed strategy to a fake Ulysses half and mocked gathers. + + Returns a factory building a fresh (sp_group, strategy) pair; each + ``pre_attention`` call consumes one key-chunks and one value-chunks list. + ``joint_fields`` picks which joint tensors the deferred Ulysses records: + the default records all three, ``("key", "value")`` mimics cached-KV reuse + (empty joint query), and ``()`` records none. + """ + head_dim = 1 + + def make(): + key_chunks = [torch.full((1, region, heads_local, head_dim), 10.0 + i) for i in range(2)] + value_chunks = [torch.full((1, region, heads_local, head_dim), 20.0 + i) for i in range(2)] + sp_group = _MockComposedSPGroup(rank=rank, allgather_world_size=2, gather_chunks=[key_chunks, value_chunks]) + return sp_group, UlyssesAllGatherKVParallelAttention(sp_group, scatter_idx=2, gather_idx=1, use_sync=False) + + def fake_pre_attention(self, query, key, value, attn_metadata, defer_joint=False): + if attn_metadata is not None and joint_fields: + joint = torch.ones(1, joint_len, heads_local, head_dim) + if "query" in joint_fields: + attn_metadata.joint_query = joint + if "key" in joint_fields: + attn_metadata.joint_key = joint + if "value" in joint_fields: + attn_metadata.joint_value = joint + return ( + torch.zeros(1, region, heads_local, head_dim), + torch.full((1, region, heads_local, head_dim), 10.0 + rank), + torch.full((1, region, heads_local, head_dim), 20.0 + rank), + attn_metadata, + "ulysses-ctx", + ) + + monkeypatch.setattr(UlyssesParallelAttention, "pre_attention", fake_pre_attention) + monkeypatch.setattr(UlyssesParallelAttention, "post_attention", lambda self, attn_output, ctx: attn_output) + return make + + +def test_composed_strategy_merges_joint_only_mask(monkeypatch): + """A joint-only 2D mask (text padding, ``attn_mask=None``) must not be dropped. + + ``defer_joint=True`` skips the Ulysses mask merge and the AllGather path + only reads ``attn_mask``: without the post-gather merge, padded joint K/V + (e.g. Qwen-Image with unequal prompt lengths) would attend unmasked. The + padding column must survive at the joint key offset of the merged mask. + """ + region, joint_len = 4, 1 + make = _composed_mask_mocks(monkeypatch, region=region, joint_len=joint_len) + query = torch.zeros(1, 2, 4, 1) + + # front: joint keys precede the gathered image keys. + joint_mask = torch.tensor([[False]], dtype=torch.bool) # one padded text token + _, strategy = make() + _, k_out, _, merged, _ = strategy.pre_attention( + query, query, query, AttentionMetadata(attn_mask=None, joint_attn_mask=joint_mask, joint_strategy="front") + ) + + full_len = joint_len + 2 * region + assert merged.attn_mask is not None, "joint-only mask was silently dropped" + assert merged.attn_mask.shape == (1, full_len) + assert merged.attn_mask.dtype == torch.bool + assert not merged.attn_mask[0, 0], "padding column must stay masked at the front joint offset" + assert bool(merged.attn_mask[0, 1:].all()), "gathered image keys must stay attendable" + assert k_out.shape[1] == full_len + + # rear: joint keys follow the gathered image keys. + _, strategy = make() + _, _, _, merged_rear, _ = strategy.pre_attention( + query, query, query, AttentionMetadata(attn_mask=None, joint_attn_mask=joint_mask, joint_strategy="rear") + ) + assert not merged_rear.attn_mask[0, -1], "padding column must stay masked at the rear joint offset" + assert bool(merged_rear.attn_mask[0, :-1].all()) + + +def test_composed_strategy_merges_joint_and_image_masks(monkeypatch): + """Both 2D masks merge in ``joint_strategy`` order, preserving both contents.""" + region, joint_len = 4, 1 + make = _composed_mask_mocks(monkeypatch, region=region, joint_len=joint_len) + query = torch.zeros(1, 2, 4, 1) + joint_mask = torch.tensor([[False]]) # one padded text token + img_mask = torch.ones(1, 2 * region, dtype=torch.bool) + img_mask[0, 3] = False # one masked image key + + _, strategy = make() + _, _, _, merged_front, _ = strategy.pre_attention( + query, + query, + query, + AttentionMetadata(attn_mask=img_mask.clone(), joint_attn_mask=joint_mask.clone(), joint_strategy="front"), + ) + assert merged_front.attn_mask.shape == (1, joint_len + 2 * region) + torch.testing.assert_close(merged_front.attn_mask, torch.cat([joint_mask, img_mask], dim=1)) + + _, strategy = make() + _, _, _, merged_rear, _ = strategy.pre_attention( + query, + query, + query, + AttentionMetadata(attn_mask=img_mask.clone(), joint_attn_mask=joint_mask.clone(), joint_strategy="rear"), + ) + torch.testing.assert_close(merged_rear.attn_mask, torch.cat([img_mask, joint_mask], dim=1)) + + +def test_composed_strategy_sizes_joint_mask_from_joint_key(monkeypatch): + """Cached-KV reuse: ``joint_query`` may be empty, but joint K/V still gather. + + The joint mask segment must be sized from ``joint_key`` (1 key here), not + from the absent joint query, and the image side must be True-filled to the + gathered global image length. + """ + region, joint_len = 4, 1 + make = _composed_mask_mocks(monkeypatch, region=region, joint_len=joint_len, joint_fields=("key", "value")) + query = torch.zeros(1, 2, 4, 1) + joint_mask = torch.tensor([[False]]) + + _, strategy = make() + q_out, k_out, _, merged, _ = strategy.pre_attention( + query, query, query, AttentionMetadata(attn_mask=None, joint_attn_mask=joint_mask, joint_strategy="front") + ) + assert q_out.shape[1] == region, "no joint query rows when joint_query is absent" + assert k_out.shape[1] == joint_len + 2 * region, "joint keys still participate as keys" + assert merged.attn_mask is not None + assert merged.attn_mask.shape == (1, joint_len + 2 * region) + assert not merged.attn_mask[0, 0] + assert bool(merged.attn_mask[0, 1:].all()) + + +def test_composed_strategy_rejects_inconsistent_2d_masks(monkeypatch): + """Shape mismatches fail closed with the expected layout in the message.""" + region, joint_len = 4, 1 + make = _composed_mask_mocks(monkeypatch, region=region, joint_len=joint_len) + query = torch.zeros(1, 2, 4, 1) + + # Image mask shorter than the gathered global image keys (e.g. a stale + # shard-local mask covering only this rank's own region). + _, strategy = make() + with pytest.raises(ValueError, match="global image keys"): + strategy.pre_attention( + query, query, query, AttentionMetadata(attn_mask=torch.ones(1, region, dtype=torch.bool)) + ) + + # Joint mask length must match the joint keys. + _, strategy = make() + with pytest.raises(ValueError, match="joint_attn_mask inconsistent"): + strategy.pre_attention( + query, + query, + query, + AttentionMetadata(attn_mask=None, joint_attn_mask=torch.ones(1, joint_len + 1, dtype=torch.bool)), + ) + + # Batch dimension must match the keys. + _, strategy = make() + with pytest.raises(ValueError, match="global image keys"): + strategy.pre_attention( + query, query, query, AttentionMetadata(attn_mask=torch.ones(2, 2 * region, dtype=torch.bool)) + ) + + +def test_composed_strategy_rejects_joint_mask_without_joint_kv(monkeypatch): + """A joint mask without joint K/V describes keys that never gather.""" + make = _composed_mask_mocks(monkeypatch, region=4, joint_len=1, joint_fields=()) + _, strategy = make() + query = torch.zeros(1, 2, 4, 1) + metadata = AttentionMetadata(attn_mask=None, joint_attn_mask=torch.ones(1, 1, dtype=torch.bool)) + + with pytest.raises(ValueError, match="without joint K/V"): + strategy.pre_attention(query, query, query, metadata) + + +def test_composed_strategy_rejects_4d_mask_with_joint_mask(monkeypatch): + """A 2D joint mask cannot be merged into a 4D image mask; fail closed.""" + region, joint_len = 4, 1 + make = _composed_mask_mocks(monkeypatch, region=region, joint_len=joint_len) + _, strategy = make() + query = torch.zeros(1, 2, 4, 1) + local_q, full_k = joint_len + region, joint_len + 2 * region + metadata = AttentionMetadata( + attn_mask=torch.ones(1, 1, local_q, full_k, dtype=torch.bool), + joint_attn_mask=torch.ones(1, joint_len, dtype=torch.bool), + joint_strategy="front", + ) + + with pytest.raises(NotImplementedError, match="cannot merge a 2D joint_attn_mask"): + strategy.pre_attention(query, query, query, metadata) + + +def test_composed_strategy_skips_region_len_collective_with_auto_pad(monkeypatch): + """advanced_uaa + framework-managed auto_pad needs no length collective.""" + make = _composed_mask_mocks(monkeypatch) + _, strategy = make() + + collective_calls: list[int] = [] + + def fake_all_gather(gathered, local, group=None): + collective_calls.append(1) + for tensor in gathered: + tensor.fill_(4) + + monkeypatch.setattr(torch.distributed, "all_gather", fake_all_gather) + monkeypatch.setattr( + "vllm_omni.diffusion.attention.parallel.ulysses_allgather.get_ulysses_mode", + lambda *, default="strict": "advanced_uaa", + ) + monkeypatch.setattr( + "vllm_omni.diffusion.attention.parallel.ulysses_allgather.is_forward_context_available", + lambda: True, + ) + + ctx_equal = SimpleNamespace(sp_rank_local_seq_lens_equal=True) + monkeypatch.setattr( + "vllm_omni.diffusion.attention.parallel.ulysses_allgather.get_forward_context", + lambda: ctx_equal, + ) + strategy._assert_equal_region_lengths(4, torch.device("cpu")) + assert collective_calls == [], "auto_pad contract makes the per-forward collective redundant" + + ctx_unequal = SimpleNamespace(sp_rank_local_seq_lens_equal=False) + monkeypatch.setattr( + "vllm_omni.diffusion.attention.parallel.ulysses_allgather.get_forward_context", + lambda: ctx_unequal, + ) + strategy._assert_equal_region_lengths(4, torch.device("cpu")) + assert collective_calls == [1], "manual sharding keeps the per-forward check" + + +def test_composed_strategy_detects_uneven_allgather_regions(monkeypatch): + """Uneven regions fail closed, and the check is skipped in strict mode.""" + sp_group = _MockComposedSPGroup( + rank=0, + allgather_world_size=2, + gather_chunks=[[torch.zeros(1, 2, 1, 1)] * 2, [torch.zeros(1, 2, 1, 1)] * 2], + ) + strategy = UlyssesAllGatherKVParallelAttention(sp_group, scatter_idx=2, gather_idx=1, use_sync=False) + + collective_calls: list[int] = [] + + def fake_all_gather(gathered, local, group=None): + collective_calls.append(1) + for tensor in gathered: + tensor.fill_(4) + + monkeypatch.setattr(torch.distributed, "all_gather", fake_all_gather) + + # Strict mode (default without a forward context): the Ulysses all-to-all + # already guarantees equal region lengths, so the check costs no collective. + strategy._assert_equal_region_lengths(4, torch.device("cpu")) + strategy._assert_equal_region_lengths(4, torch.device("cpu")) + assert collective_calls == [] + + # advanced_uaa: rank-local region lengths may legitimately differ, so the + # check runs on every call and fails closed when they do. + monkeypatch.setattr( + "vllm_omni.diffusion.attention.parallel.ulysses_allgather.get_ulysses_mode", + lambda *, default="strict": "advanced_uaa", + ) + uaa_strategy = UlyssesAllGatherKVParallelAttention(sp_group, scatter_idx=2, gather_idx=1, use_sync=False) + + def fake_all_gather_uneven(gathered, local, group=None): + gathered[0].fill_(4) + gathered[1].fill_(5) + + monkeypatch.setattr(torch.distributed, "all_gather", fake_all_gather_uneven) + with pytest.raises(ValueError, match="equally long region"): + uaa_strategy._assert_equal_region_lengths(4, torch.device("cpu")) + + @hardware_test(res={"cuda": "L4"}, num_cards=4) @pytest.mark.parametrize( "test_model_cls", @@ -422,6 +976,7 @@ def test_allgather_kv_keeps_gathered_kv_compressed_for_gqa(): [ pytest.param(2, 2, 1, None, id="ulysses-ring"), pytest.param(1, 1, 2, None, id="allgather-kv"), + pytest.param(2, 1, 2, None, id="ulysses-allgather-kv"), pytest.param(1, 2, 1, 2, id="ring-gqa"), pytest.param(1, 2, 1, 1, id="ring-mqa"), ], @@ -451,8 +1006,8 @@ def test_sequence_parallel( head_size: int, num_kv_heads: int | None, ): - """Compare Ulysses/Ring and AllGather-KV SP against a single-rank run.""" - sequence_parallel_size = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree + """Compare Ulysses/Ring, AllGather-KV and Ulysses x AllGather-KV SP against a single-rank run.""" + sequence_parallel_size = ulysses_degree * ring_degree * allgather_degree # Skip if not enough GPUs available available_gpus = current_omni_platform.get_device_count() @@ -460,14 +1015,14 @@ def test_sequence_parallel( pytest.skip(f"Test requires {sequence_parallel_size} GPUs but only {available_gpus} available") # Create temporary files to share results between processes - with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: - baseline_output_file = f.name - with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: - sp_output_file = f.name - with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: - model_state_file = f.name - with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: - input_data_file = f.name + with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as tmp_file: + baseline_output_file = tmp_file.name + with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as tmp_file: + sp_output_file = tmp_file.name + with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as tmp_file: + model_state_file = tmp_file.name + with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as tmp_file: + input_data_file = tmp_file.name try: # Step 1: Run without SP (baseline with ulysses_degree=1, ring_degree=1) @@ -534,8 +1089,8 @@ def test_sequence_parallel( # Step 3: Verify input consistency and compare outputs print(f"\n{'=' * 80}") print("Verifying input data consistency...") - with open(input_data_file, "rb") as f: - input_data = pickle.load(f) + with open(input_data_file, "rb") as fh: + input_data = pickle.load(fh) input_checksum = hash(input_data.tobytes()) print(f" Input data shape: {input_data.shape}") print(f" Input data checksum: {input_checksum}") @@ -543,10 +1098,10 @@ def test_sequence_parallel( print(f"\n{'=' * 80}") print("Comparing outputs between baseline and SP...") - with open(baseline_output_file, "rb") as f: - baseline_output = pickle.load(f) - with open(sp_output_file, "rb") as f: - sp_output = pickle.load(f) + with open(baseline_output_file, "rb") as fh: + baseline_output = pickle.load(fh) + with open(sp_output_file, "rb") as fh: + sp_output = pickle.load(fh) # Convert to tensors for comparison baseline_tensor = torch.tensor(baseline_output) @@ -597,9 +1152,9 @@ def test_sequence_parallel( finally: # Clean up temporary files - for f in [baseline_output_file, sp_output_file, model_state_file, input_data_file]: - if os.path.exists(f): - os.remove(f) + for path in [baseline_output_file, sp_output_file, model_state_file, input_data_file]: + if os.path.exists(path): + os.remove(path) def ulysses_attention_on_test_model( @@ -630,10 +1185,7 @@ def ulysses_attention_on_test_model( RANDOM_SEED = 42 seed_everything(RANDOM_SEED) - if allgather_degree > 1: - sp_kind = f"allgather={allgather_degree}" - else: - sp_kind = f"ulysses={ulysses_degree}, ring={ring_degree}" + sp_kind = _sp_kind_label(ulysses_degree, ring_degree, allgather_degree) mode_str = "Baseline (no SP)" if is_baseline else f"SP ({sp_kind})" print(f"\n[{mode_str}] Rank {local_rank}/{world_size} - Random seed set to {RANDOM_SEED}") @@ -802,10 +1354,7 @@ def ulysses_attention_on_test_model( with open(output_file, "wb") as f: pickle.dump(output_np, f) - if allgather_degree > 1: - sp_kind = f"allgather={allgather_degree}" - else: - sp_kind = f"ulysses={ulysses_degree}, ring={ring_degree}" + sp_kind = _sp_kind_label(ulysses_degree, ring_degree, allgather_degree) mode_str = "baseline (no SP)" if is_baseline else f"SP ({sp_kind})" print( f"\n[{mode_str}] ✓ Saved output with shape {full_output.shape}:\n" @@ -815,3 +1364,154 @@ def ulysses_attention_on_test_model( ) destroy_distributed_env() + + +def _packed_fa_worker( + local_rank: int, + world_size: int, + docs: list[int], + num_heads: int, + head_size: int, + dtype: torch.dtype, + output_file: str, +): + """Run packed FlashAttention varlen with or without AllGather-KV SP. + + Every rank generates the identical global packed input from the same + seed, then feeds only its contiguous query shard (with the *global* + packed metadata, as a producer would) to the production Attention layer. + Under SP the AllGather-KV strategy gathers global K/V and must remap the + packed query boundaries per rank before the real CUDA varlen kernel. + """ + seed_everything(42) + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) + torch.set_default_device(device) + torch.set_default_dtype(dtype) + + update_environment_variables( + { + "RANK": str(local_rank), + "LOCAL_RANK": str(local_rank), + "WORLD_SIZE": str(world_size), + "MASTER_ADDR": "localhost", + "MASTER_PORT": "12347", + } + ) + init_distributed_environment() + + total = docs[-1] + local_len = total // world_size + parallel_config = DiffusionParallelConfig( + pipeline_parallel_size=1, + data_parallel_size=1, + tensor_parallel_size=1, + sequence_parallel_size=world_size, + ulysses_degree=1, + ring_degree=1, + allgather_degree=world_size, + cfg_parallel_size=1, + ) + od_config = OmniDiffusionConfig.from_kwargs( + model="test_model", + dtype=dtype, + parallel_config=parallel_config, + diffusion_attention_backend="FLASH_ATTN", + ) + initialize_model_parallel( + data_parallel_size=1, + cfg_parallel_size=1, + sequence_parallel_size=world_size, + ulysses_degree=1, + ring_degree=1, + allgather_degree=world_size, + tensor_parallel_size=1, + pipeline_parallel_size=1, + ) + + with set_forward_context(omni_diffusion_config=od_config), set_current_diffusion_config(od_config): + attn = Attention( + num_heads=num_heads, + head_size=head_size, + causal=False, + softmax_scale=1.0 / (head_size**0.5), + scatter_idx=2, + gather_idx=1, + use_sync=False, + ) + q = torch.randn(1, total, num_heads, head_size) + k = torch.randn(1, total, num_heads, head_size) + v = torch.randn(1, total, num_heads, head_size) + max_len = max(b - a for a, b in zip(docs[:-1], docs[1:])) + cu = torch.tensor(docs, dtype=torch.int32, device=device) + metadata = AttentionMetadata( + extra={ + "cu_seqlens_q": cu, + "cu_seqlens_k": cu, + "max_seqlen_q": max_len, + "max_seqlen_k": max_len, + } + ) + start = local_rank * local_len + out = attn( + q[:, start : start + local_len], + k[:, start : start + local_len], + v[:, start : start + local_len], + metadata, + ) + with open(f"{output_file}.{local_rank}", "wb") as f: + pickle.dump(out.detach().cpu().float().numpy(), f) + + destroy_distributed_env() + + +def test_packed_allgather_matches_unsplit_flash_varlen(): + """Real CUDA varlen kernel: AllGather-KV SP output == unsplit packed attention. + + Two packed requests (8 + 4 tokens, boundaries [0, 8, 12]) with + allgather_degree=2: document 0 spans the shard boundary, so rank 0 keeps + 6 of its query rows and rank 1 keeps 2 plus all 4 rows of request 1 + (local cu_seqlens_q [0, 6, 6] / [0, 2, 6]). The baseline is the same + FlashAttention varlen kernel on the unsplit sequence with the global + boundaries. This pins the review's reproduction end-to-end: request + isolation, boundary legality, and numerical agreement. + """ + if current_omni_platform.get_device_count() < 2: + pytest.skip("Test requires 2 GPUs but only fewer are available") + + from vllm_omni.diffusion.attention.backends.utils.fa import flash_attn_varlen_func + + if flash_attn_varlen_func is None: + pytest.skip("flash_attn_varlen_func is unavailable") + + docs = [0, 8, 12] + num_heads, head_size = 8, 64 + + with tempfile.TemporaryDirectory() as tmpdir: + base_file = os.path.join(tmpdir, "baseline.pkl") + sp_file = os.path.join(tmpdir, "sp.pkl") + + # Baseline: single rank, no SP, global packed metadata as-is. + torch.multiprocessing.spawn( + _packed_fa_worker, + args=(1, docs, num_heads, head_size, torch.bfloat16, base_file), + nprocs=1, + ) + # SP: two ranks with allgather_degree=2, each feeding its shard with + # the producer's *global* packed metadata. + torch.multiprocessing.spawn( + _packed_fa_worker, + args=(2, docs, num_heads, head_size, torch.bfloat16, sp_file), + nprocs=2, + ) + + with open(base_file + ".0", "rb") as fh: + baseline = torch.tensor(pickle.load(fh)) + sp_rows = [] + for rank in range(2): + with open(f"{sp_file}.{rank}", "rb") as fh: + sp_rows.append(torch.tensor(pickle.load(fh))) + sp_output = torch.cat(sp_rows, dim=1) + + assert sp_output.shape == baseline.shape + torch.testing.assert_close(sp_output, baseline, atol=5e-2, rtol=5e-2) diff --git a/tests/diffusion/distributed/test_parallel_state_sp_groups.py b/tests/diffusion/distributed/test_parallel_state_sp_groups.py index 3aaefbcf94f..8081ddf855a 100644 --- a/tests/diffusion/distributed/test_parallel_state_sp_groups.py +++ b/tests/diffusion/distributed/test_parallel_state_sp_groups.py @@ -12,7 +12,11 @@ from vllm.utils.network_utils import get_file_store_init_method from vllm_omni.diffusion.distributed import parallel_state as omni_parallel_state -from vllm_omni.diffusion.distributed.parallel_state import RankGenerator, set_seq_parallel_pg +from vllm_omni.diffusion.distributed.parallel_state import ( + RankGenerator, + build_ulysses_allgather_rank_groups, + set_seq_parallel_pg, +) pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.core_model, pytest.mark.cpu] @@ -118,17 +122,105 @@ def test_set_seq_parallel_pg_builds_independent_allgather_group(rank, monkeypatc @pytest.mark.cpu -def test_set_seq_parallel_pg_rejects_allgather_with_ulysses_or_ring(monkeypatch): +@pytest.mark.parametrize( + "rank, expected_ulysses, expected_allgather, expected_ring", + [ + (0, [0, 1], [0, 2], [0]), + (1, [0, 1], [1, 3], [1]), + (2, [2, 3], [0, 2], [2]), + (3, [2, 3], [1, 3], [3]), + ], +) +def test_set_seq_parallel_pg_builds_orthogonal_ulysses_allgather_groups( + rank, expected_ulysses, expected_allgather, expected_ring, monkeypatch +): + """Ulysses x AllGather-KV: contiguous Ulysses blocks + stride-U AllGather groups.""" + created_groups: list[SimpleNamespace] = [] + monkeypatch.setattr(torch.distributed, "new_group", _fake_new_group_factory(created_groups)) + + ulysses_pg, ring_pg, allgather_pg = set_seq_parallel_pg( + sp_ulysses_degree=2, + sp_ring_degree=1, + sp_allgather_degree=2, + rank=rank, + world_size=4, + sp_group_ranks=[[0, 1, 2, 3]], + ) + + assert ulysses_pg.ranks == expected_ulysses + assert allgather_pg.ranks == expected_allgather + assert ring_pg.ranks == expected_ring + + +@pytest.mark.cpu +@pytest.mark.parametrize(("ulysses", "allgather"), [(2, 2), (4, 2), (2, 4), (4, 4), (1, 4), (4, 1)]) +def test_build_ulysses_allgather_rank_groups_is_orthogonal(ulysses: int, allgather: int): + """Every Ulysses group must intersect every AllGather group in exactly one rank. + + That orthogonality plus "Ulysses groups are contiguous" is what makes the + composed transform reproduce the single-rank sequence order. + """ + sp_size = ulysses * allgather + sp_group_ranks = [list(range(sp_size))] + + ulysses_groups, allgather_groups = build_ulysses_allgather_rank_groups(sp_group_ranks, ulysses, allgather) + + assert len(ulysses_groups) == allgather + assert len(allgather_groups) == ulysses + assert all(len(g) == ulysses for g in ulysses_groups) + assert all(len(g) == allgather for g in allgather_groups) + + # Orthogonal and covering: each Ulysses group meets each AllGather group in + # exactly one rank, and the Ulysses groups tile the SP group in order. + assert sorted(r for g in ulysses_groups for r in g) == list(range(sp_size)) + assert ulysses_groups == [list(range(i * ulysses, (i + 1) * ulysses)) for i in range(allgather)] + for u_group in ulysses_groups: + for a_group in allgather_groups: + assert len(set(u_group) & set(a_group)) == 1 + + # allgather_rank is the index inside the AllGather group, and it must match + # p // U so that concatenating a region per allgather_rank rebuilds the + # global sequence in order. + for a_group in allgather_groups: + for position, p in enumerate(a_group): + assert p // ulysses == position + + +@pytest.mark.cpu +@pytest.mark.parametrize("rank", [0, 1, 2, 3, 4, 5, 6, 7]) +def test_orthogonal_groups_span_non_contiguous_sp_group(rank, monkeypatch): + """SP groups need not be contiguous in global rank (tp=2, sp=4 -> [0,2,4,6]).""" + created_groups: list[SimpleNamespace] = [] + monkeypatch.setattr(torch.distributed, "new_group", _fake_new_group_factory(created_groups)) + + ulysses_pg, _, allgather_pg = set_seq_parallel_pg( + sp_ulysses_degree=2, + sp_ring_degree=1, + sp_allgather_degree=2, + rank=rank, + world_size=8, + sp_group_ranks=[[0, 2, 4, 6], [1, 3, 5, 7]], + ) + + offset = rank % 2 + local = [offset, offset + 2, offset + 4, offset + 6] + position = local.index(rank) + assert ulysses_pg.ranks == local[(position // 2) * 2 : (position // 2 + 1) * 2] + assert allgather_pg.ranks == local[position % 2 :: 2] + + +@pytest.mark.cpu +def test_set_seq_parallel_pg_rejects_allgather_composed_with_ring(monkeypatch): created_groups: list[SimpleNamespace] = [] monkeypatch.setattr(torch.distributed, "new_group", _fake_new_group_factory(created_groups)) - with pytest.raises(ValueError, match="mutually exclusive"): + with pytest.raises(ValueError, match="cannot be composed with Ring"): set_seq_parallel_pg( sp_ulysses_degree=2, - sp_ring_degree=1, + sp_ring_degree=2, sp_allgather_degree=2, rank=0, - world_size=4, + world_size=8, ) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index d581309e359..413a34a9b85 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -714,11 +714,13 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): def __post_init__(self) -> None: self.data_parallel_index = self.data_parallel_rank - self.sequence_parallel_size = ( - self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree - ) - if self.allgather_degree > 1 and (self.ulysses_degree > 1 or self.ring_degree > 1): - raise ValueError("allgather_degree > 1 is mutually exclusive with ulysses_degree/ring_degree > 1") + self.sequence_parallel_size = self.ulysses_degree * self.ring_degree * self.allgather_degree + if self.allgather_degree > 1 and self.ring_degree > 1: + raise ValueError( + "allgather_degree > 1 cannot be composed with ring_degree > 1: the supported " + "two-dimensional diffusion SP topology is Ulysses x AllGather-KV " + "(ulysses_degree > 1 together with allgather_degree > 1, ring_degree == 1)." + ) if self.ulysses_mode not in {"strict", "advanced_uaa"}: raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'") if self.vae_parallel_mode not in {"tile", "batch", "spatial_shard_height", "spatial_shard_width"}: diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index b2a588ac0cd..2b59f4ed153 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -154,9 +154,7 @@ def _apply_diffusion_parallel_runtime_overrides( ulysses_degree = parallel_config_dict.get("ulysses_degree") or 1 ring_degree = parallel_config_dict.get("ring_degree") or 1 allgather_degree = parallel_config_dict.get("allgather_degree") or 1 - parallel_config_dict["sequence_parallel_size"] = ( - allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree - ) + parallel_config_dict["sequence_parallel_size"] = ulysses_degree * ring_degree * allgather_degree if parallel_config_dict is not None: engine_args["parallel_config"] = parallel_config_dict diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index d239525c245..a7e724fcfad 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -197,7 +197,7 @@ def __init__( self.backend_pref = attn_backend_cls.get_name() logger.debug("Attention(role=%s) → platform default (%s)", role, self.backend_pref) - self.attn_backend: type[AttentionBackend] | None = attn_backend_cls + self.attn_backend = attn_backend_cls self.attn_impl_cls = self.attn_backend.get_impl_cls() if impl_overrides is not None: override = impl_overrides.get(attn_backend_cls.get_name()) @@ -223,7 +223,7 @@ def __init__( ) # Compatibility kernels run inside shared dispatch, between the # parallel strategy's input preparation and output restoration. - self.sdpa_fallback: AttentionImpl | None = SDPABackend.get_impl_cls()( + self.sdpa_fallback = SDPABackend.get_impl_cls()( num_heads=num_heads, head_size=head_size, softmax_scale=softmax_scale, @@ -496,8 +496,10 @@ def _forward_impl( raise NotImplementedError( "paged Scheduler KV is not supported with Ring attention; use strict Ulysses or no SP" ) - if strategy_name == "allgather_kv": - raise NotImplementedError("paged Scheduler KV is not supported with AllGather-KV sequence parallelism") + if strategy_name in ("allgather_kv", "ulysses_allgather_kv"): + raise NotImplementedError( + f"paged Scheduler KV is not supported with the {strategy_name!r} sequence-parallel strategy" + ) if strategy_name == "ulysses" and get_ulysses_mode(default="strict") != "strict": raise NotImplementedError("paged Scheduler KV currently supports only strict Ulysses") diff --git a/vllm_omni/diffusion/attention/parallel/allgather_kv.py b/vllm_omni/diffusion/attention/parallel/allgather_kv.py index 51a463bb19e..f5ff0ff3f12 100644 --- a/vllm_omni/diffusion/attention/parallel/allgather_kv.py +++ b/vllm_omni/diffusion/attention/parallel/allgather_kv.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations from dataclasses import dataclass, replace -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch @@ -60,8 +60,8 @@ def pre_attention( if joint_strategy not in {"front", "rear"}: raise ValueError(f"Unsupported joint_strategy: {joint_strategy!r}") - k_img_full = self._sp_group.all_gather(key, dim=1, group=self._allgather_group) - v_img_full = self._sp_group.all_gather(value, dim=1, group=self._allgather_group) + k_img_full: torch.Tensor = self._sp_group.all_gather(key, dim=1, group=self._allgather_group) + v_img_full: torch.Tensor = self._sp_group.all_gather(value, dim=1, group=self._allgather_group) if joint_k is not None: if joint_k.shape[2] != key.shape[2]: @@ -108,6 +108,67 @@ def pre_attention( return q_local, k_full, v_full, attn_metadata, _AllGatherKVCtx(name=self.name) + def _remap_packed_query_metadata( + self, + attn_metadata: AttentionMetadata, + *, + img_start: int, + img_seq_local: int, + img_seq_full: int, + joint_len: int, + ) -> dict[str, Any]: + """Rebuild packed query boundaries for this rank's local query shard. + + The producer's ``extra["cu_seqlens_q"]`` / ``max_seqlen_q`` describe + the *global* packed layout, but after the gather this rank's Q holds + only the image rows ``[img_start, img_start + img_seq_local)``. A + packed backend (e.g. CUDA FlashAttention varlen) consumes those + boundaries directly and never reads ``query_ranges``, so they must be + intersected with the local span here. Every global document keeps + exactly one -- possibly zero-length -- local entry: dropping empty + entries would break the one-to-one pairing with the global key + documents in ``cu_seqlens_k``. + + The intersection runs as device tensor ops, so no attention layer + pays a ``.item()`` host sync. Key-side metadata stays global (K/V are + gathered); ``max_seqlen_q`` keeps the producer value, which remains a + safe upper bound of every local segment (kernels only use it for + scheduling). ``packed_padding`` (the canonical single-document + ``[0, total]`` view) narrows to the local valid query prefix while + its key length keeps the global valid prefix, so consumers slice Q + correctly and Q≠KV-length consumers fall back instead of truncating + the gathered K/V. + + Packed metadata combined with joint attention is not defined yet (no + producer emits both, and joint K/V would shift the key document + boundaries); it fails closed. + """ + extra = attn_metadata.extra + if not extra or "cu_seqlens_q" not in extra: + return {} + if joint_len or attn_metadata.joint_key is not None: + raise NotImplementedError( + "AllGather-KV SP does not yet support packed cu_seqlens_q combined with joint " + "attention. Run one request per forward, or set allgather_degree=1." + ) + cu_q = extra["cu_seqlens_q"] + span_lo, span_hi = img_start, img_start + img_seq_local + local_lens = torch.clamp(torch.clamp(cu_q[1:], max=span_hi) - torch.clamp(cu_q[:-1], min=span_lo), min=0) + local_cu = torch.zeros_like(cu_q) + local_cu[1:] = torch.cumsum(local_lens, dim=0).to(cu_q.dtype) + + updates: dict[str, Any] = {"extra": {**extra, "cu_seqlens_q": local_cu}} + packed_padding = attn_metadata.packed_padding + if packed_padding is not None: + valid_kv = int(extra.get("valid_kv_length") or img_seq_full) + local_valid = max(min(valid_kv, span_hi) - span_lo, 0) + updates["packed_padding"] = replace( + packed_padding, + q_length=local_valid, + cu_seqlens_q=local_cu[:2].clone(), + ) + return updates + def _slice_attn_metadata_for_local_query( self, attn_metadata: AttentionMetadata | None, @@ -142,12 +203,20 @@ def _slice_attn_metadata_for_local_query( *((QueryRange(img_seq_local, q_local_len, query_global_offset + img_seq_full),) if joint_len else ()), ) + packed_updates = self._remap_packed_query_metadata( + attn_metadata, + img_start=img_start, + img_seq_local=img_seq_local, + img_seq_full=img_seq_full, + joint_len=joint_len, + ) + if attn_metadata.attn_mask is None: - return replace(attn_metadata, query_ranges=ranges) + return replace(attn_metadata, query_ranges=ranges, **packed_updates) mask = attn_metadata.attn_mask if mask.ndim != 4: - return replace(attn_metadata, query_ranges=ranges) + return replace(attn_metadata, query_ranges=ranges, **packed_updates) if mask.shape[-2] == q_local_len: local_mask = mask @@ -157,7 +226,7 @@ def _slice_attn_metadata_for_local_query( raise ValueError( "AllGather-KV SP received an attention mask with incompatible Q length: " f"mask_q={mask.shape[-2]}, expected local_q={q_local_len} or full_q={q_full_len} " - f"(joint_len={joint_len}, img_seq_local={img_seq_local}, img_seq_full={img_seq_full})." + f"(joint_len={joint_len}, img_seq_local={img_seq_local}, img_seq_full={img_seq_full}." ) parts = [ @@ -172,7 +241,7 @@ def _slice_attn_metadata_for_local_query( "AllGather-KV SP produced an attention mask with incompatible local Q length: " f"mask_q={local_mask.shape[-2]}, q_local={q_local_len}." ) - return replace(attn_metadata, attn_mask=local_mask.contiguous(), query_ranges=ranges) + return replace(attn_metadata, attn_mask=local_mask.contiguous(), query_ranges=ranges, **packed_updates) def post_attention( self, diff --git a/vllm_omni/diffusion/attention/parallel/factory.py b/vllm_omni/diffusion/attention/parallel/factory.py index a7e41575b1b..68fa96b7232 100644 --- a/vllm_omni/diffusion/attention/parallel/factory.py +++ b/vllm_omni/diffusion/attention/parallel/factory.py @@ -14,6 +14,9 @@ ) from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention +from vllm_omni.diffusion.attention.parallel.ulysses_allgather import ( + UlyssesAllGatherKVParallelAttention, +) from vllm_omni.diffusion.distributed.parallel_state import ( get_sequence_parallel_world_size, get_sp_group, @@ -40,6 +43,8 @@ def build_parallel_attention_strategy( if not is_forward_context_available(): return NoParallelAttention() cfg = get_forward_context().omni_diffusion_config + if cfg is None: + return NoParallelAttention() p = cfg.parallel_config ulysses_degree = getattr(p, "ulysses_degree", 1) @@ -64,15 +69,31 @@ def build_parallel_attention_strategy( f"allgather={allgather_degree}), but the initialized SP world size is not greater than one." ) + # AllGather-KV, optionally composed with Ulysses as an orthogonal 2D + # sequence-parallel topology. K/V are gathered over the AllGather group + # *after* any Ulysses reshard, so the attention backend always runs a + # local-Q/global-KV problem. if allgather_degree > 1: if causal: raise ValueError("AllGather-KV SP only supports non-causal attention.") - if ulysses_degree > 1 or ring_degree > 1: + if ring_degree > 1: raise ValueError( - f"AllGather-KV SP is mutually exclusive with Ulysses/Ring in v1 " + "AllGather-KV SP cannot be composed with Ring " f"(got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, " f"allgather_degree={allgather_degree})." ) + if ulysses_degree > 1: + logger.debug( + f"Using UlyssesAllGatherKVParallelAttention " + f"(ulysses_degree={ulysses_degree}, allgather_degree={allgather_degree})" + ) + return UlyssesAllGatherKVParallelAttention( + sp_group=sp_group, + scatter_idx=scatter_idx, + gather_idx=gather_idx, + use_sync=use_sync, + ulysses_a2a_permute=ulysses_a2a_permute, + ) logger.debug(f"Using AllGatherKVParallelAttention (allgather_degree={allgather_degree})") return AllGatherKVParallelAttention(sp_group=sp_group) diff --git a/vllm_omni/diffusion/attention/parallel/ulysses.py b/vllm_omni/diffusion/attention/parallel/ulysses.py index 28298c0a280..95937b1a05e 100644 --- a/vllm_omni/diffusion/attention/parallel/ulysses.py +++ b/vllm_omni/diffusion/attention/parallel/ulysses.py @@ -278,7 +278,18 @@ def pre_attention( key: torch.Tensor, value: torch.Tensor, attn_metadata: AttentionMetadata | None, + defer_joint: bool = False, ): + """Reshard Q/K/V for Ulysses. + + Args: + defer_joint: Leave joint q/k/v out of the resharded tensors and + only record this rank's head-sliced joint tensors (and any 2D + key mask, untouched) in ``attn_metadata``. The composed + Ulysses x AllGather-KV strategy re-attaches them after its own + K/V gather; attaching them earlier would replicate the joint + K/V once per AllGather rank. + """ mode = get_ulysses_mode(default="strict") ulysses_world_size = self._sp_group.ulysses_world_size gate_compress = None @@ -406,6 +417,11 @@ def pre_attention( if attn_metadata is not None: attn_metadata.joint_key = joint_tensor_key attn_metadata.joint_value = joint_tensor_value + # Also expose the head-sliced joint query unconditionally so the + # metadata stays q/k/v-consistent; the non-defer paths below + # concatenate the joint tensors into Q/K/V and never read this + # field back, so writing it is harmless there. + attn_metadata.joint_query = joint_tensor_query ulysses_world_size = self._sp_group.ulysses_world_size if mode == "advanced_uaa": @@ -550,7 +566,7 @@ def pre_attention( local_seq_len = 0 orig_head_cnt = 0 - if is_joint: + if is_joint and not defer_joint: # Concatenate joint query AFTER AllToAll # Image query is now (B, S, H/P, D). Joint query is (B, S_txt, H/P, D). # This is dimensionally consistent. @@ -564,7 +580,7 @@ def pre_attention( # Instead, they should remain in attn_metadata and be passed to the Ring kernel. use_ring = self._sp_group.ring_world_size > 1 - if is_joint and not use_ring: + if is_joint and not use_ring and not defer_joint: # Concatenate joint key/value after all-to-all ONLY for pure Ulysses (Local Attention). if joint_strategy == "front": key = torch.cat([joint_tensor_key, key], dim=1) @@ -600,7 +616,7 @@ def pre_attention( if attn_metadata.joint_attn_mask is not None and attn_metadata.joint_attn_mask.ndim == 2: use_2d_mask = True - if attn_metadata is not None and use_2d_mask: + if attn_metadata is not None and use_2d_mask and not defer_joint: if is_joint: if attn_metadata.joint_attn_mask is None and attn_metadata.attn_mask is None: attn_metadata.attn_mask = None diff --git a/vllm_omni/diffusion/attention/parallel/ulysses_allgather.py b/vllm_omni/diffusion/attention/parallel/ulysses_allgather.py new file mode 100644 index 00000000000..064d07b01af --- /dev/null +++ b/vllm_omni/diffusion/attention/parallel/ulysses_allgather.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.distributed as dist + +from vllm_omni.diffusion.attention.parallel.allgather_kv import ( + AllGatherKVParallelAttention, +) +from vllm_omni.diffusion.attention.parallel.base import ParallelAttentionContext +from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention +from vllm_omni.diffusion.distributed.group_coordinator import ( + SequenceParallelGroupCoordinator, +) +from vllm_omni.diffusion.forward_context import ( + get_forward_context, + get_ulysses_mode, + is_forward_context_available, +) + +if TYPE_CHECKING: + from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata + + +class UlyssesAllGatherKVParallelAttention(AllGatherKVParallelAttention): + """Composed 2D sequence parallelism: Ulysses (heads) x AllGather-KV (sequence). + + With ``U = ulysses_degree`` and ``A = allgather_degree`` (``ring_degree`` + must be 1), each rank starts from the flat contiguous SP shard and runs:: + + Ulysses all-to-all (U group): Q, K, V -> [B, S / A, H / U, D] + K/V AllGather (A group): K, V -> [B, S, H / U, D] + local-Q / global-KV attention: O -> [B, S / A, H / U, D] + reverse Ulysses all-to-all: O -> [B, S / (U x A), H, D] + + The kernel sees the same local-Q/global-KV problem as plain AllGather-KV + (stable global K/V indexing, unlike Ring). The collective order and the + rank layout it relies on are load-bearing; see + ``build_ulysses_allgather_rank_groups`` for the invariant. + + Both halves are reused rather than reimplemented: Ulysses runs with + ``defer_joint=True`` (joint q/k/v stay head-sliced in the metadata), + then the inherited AllGather-KV ``pre_attention`` gathers the image K/V and + re-attaches the joint tensors -- attaching them before the gather would + replicate them once per AllGather rank. + + Fails closed on: ``ring_degree > 1``, causal attention, a non-2D + ``joint_attn_mask``, a 2D ``joint_attn_mask`` combined with a 4D + ``attn_mask``, 2D key masks that do not cover the gathered global image + keys, and uneven per-rank region lengths. All 2D key-mask merging (image + and joint) is owned by ``_merge_joint_attn_mask`` against the final, + post-gather key layout; 4D masks keep the inherited AllGather query-range + slicing. + """ + + def __init__( + self, + sp_group: SequenceParallelGroupCoordinator, + scatter_idx: int, + gather_idx: int, + use_sync: bool, + ulysses_a2a_permute: bool = False, + ) -> None: + super().__init__(sp_group) + self._ulysses = UlyssesParallelAttention( + sp_group, + scatter_idx=scatter_idx, + gather_idx=gather_idx, + use_sync=use_sync, + ulysses_a2a_permute=ulysses_a2a_permute, + ) + + @property + def name(self) -> str: + return "ulysses_allgather_kv" + + def _assert_equal_region_lengths(self, region_len: int, device: torch.device) -> None: + """Fail fast when the ``A`` regions do not have equal length. + + ``all_gather_into_tensor`` derives its output shape from the *local* + input, so uneven regions would hang or corrupt instead of raising. In + strict mode the Ulysses all-to-all already guarantees equal region + lengths (seq must be evenly shardable), so no collective is spent; + under advanced_uaa rank-local lengths may legitimately differ, so the + check runs on every forward there -- unless every active SP boundary + is framework-managed auto_pad, whose contract already guarantees equal + local lengths. + + The cheap early returns live outside the ``torch.compiler.disable``d + checker so strict-mode compiled forwards do not unconditionally split + the graph to call a no-op. + """ + if get_ulysses_mode(default="strict") == "strict" or self._sp_size <= 1: + return + if is_forward_context_available() and get_forward_context().sp_rank_local_seq_lens_equal: + # auto_pad made every rank's shard equally long at each SP + # boundary, so equality is known without a per-forward collective + # or a host sync (mirrors the advanced_uaa fast path in Ulysses). + return + self._check_equal_region_lengths_collective(region_len, device) + + @torch.compiler.disable + def _check_equal_region_lengths_collective(self, region_len: int, device: torch.device) -> None: + local = torch.tensor([int(region_len)], dtype=torch.int64, device=device) + gathered = [torch.empty_like(local) for _ in range(self._sp_size)] + dist.all_gather(gathered, local, group=self._allgather_group) + lengths = [int(t.item()) for t in gathered] + if len(set(lengths)) != 1: + raise ValueError( + "Ulysses x AllGather-KV requires every AllGather rank to hold an equally long " + "region after the Ulysses all-to-all, but got region lengths " + f"{lengths} across allgather ranks. This means the shared sequence was not evenly " + "shardable across the SP group. Choose a sequence length divisible by " + "ulysses_degree * allgather_degree, or enable auto_pad in the model's _sp_plan." + ) + + @staticmethod + def _merge_joint_attn_mask(attn_metadata: AttentionMetadata | None, key: torch.Tensor): + """Own all 2D key-mask handling against the final, post-gather keys. + + ``defer_joint=True`` skips the Ulysses mask merge and the AllGather + path passes 2D masks through untouched. A 2D mask marks invalid keys, + so merging is concatenation: True-fill the missing side -- the joint + side is sized from ``joint_key`` because ``joint_query`` can be empty + during cached-KV reuse -- then concatenate in ``joint_strategy`` + order. A provided image mask must cover the gathered global image + keys, which every rank sees identically. + + 4D masks are not key-validity masks: they are left to the inherited + AllGather query-range slicing, and combined with a 2D joint mask they + fail closed instead of being merged. + """ + if attn_metadata is None: + return attn_metadata + joint_mask = attn_metadata.joint_attn_mask + img_mask = attn_metadata.attn_mask + if joint_mask is None and img_mask is None: + return attn_metadata + + # Rank check first: the merge only combines 2D masks (an absent side + # counts as 2D -- it gets True-filled below). + joint_rank = joint_mask.ndim if joint_mask is not None else 2 + img_rank = img_mask.ndim if img_mask is not None else 2 + if joint_rank != 2: + raise NotImplementedError( + f"Ulysses x AllGather-KV only supports a 2D joint_attn_mask (got ndim={joint_rank})." + ) + if img_rank == 4: + if joint_mask is not None: + raise NotImplementedError( + "Ulysses x AllGather-KV cannot merge a 2D joint_attn_mask with a 4D " + "attn_mask: move the padding into a single mask, or set allgather_degree=1." + ) + return attn_metadata + if img_rank != 2: + raise NotImplementedError(f"Ulysses x AllGather-KV only supports 2D or 4D attn_mask (got ndim={img_rank}).") + + joint_k = attn_metadata.joint_key + if joint_mask is not None and joint_k is None: + raise ValueError( + "Ulysses x AllGather-KV got a joint_attn_mask without joint K/V: " + "the mask describes joint keys that would never be gathered." + ) + + batch = key.shape[0] + joint_len = joint_k.shape[1] if joint_k is not None else 0 + img_len = key.shape[1] - joint_len + if img_mask is not None and tuple(img_mask.shape) != (batch, img_len): + raise ValueError( + "Ulysses x AllGather-KV expects a 2D attn_mask covering the gathered " + f"global image keys [batch={batch}, keys={img_len}], got {tuple(img_mask.shape)}." + ) + if joint_mask is not None and tuple(joint_mask.shape) != (batch, joint_len): + raise ValueError( + "Ulysses x AllGather-KV got a joint_attn_mask inconsistent with the joint " + f"keys: expected [batch={batch}, keys={joint_len}], got {tuple(joint_mask.shape)}." + ) + + if joint_len == 0: + # No joint keys: the image mask already covers every key. + assert img_mask is not None + attn_metadata.attn_mask = img_mask.bool().contiguous() + return attn_metadata + + if joint_mask is None: + joint_mask = torch.ones(batch, joint_len, dtype=torch.bool, device=key.device) + if img_mask is None: + img_mask = torch.ones(batch, img_len, dtype=torch.bool, device=key.device) + joint_strategy = attn_metadata.joint_strategy or "front" + merged = ( + torch.cat([joint_mask, img_mask], dim=1) + if joint_strategy == "front" + else torch.cat([img_mask, joint_mask], dim=1) + ) + attn_metadata.attn_mask = merged.bool().contiguous() + return attn_metadata + + def pre_attention( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None, + ): + # 1. Ulysses: reshard the image shard to [B, S/A, H/U, D] and record the + # head-sliced joint tensors for later, without concatenating them. + query, key, value, attn_metadata, ctx = self._ulysses.pre_attention( + query, + key, + value, + attn_metadata, + defer_joint=True, + ) + + # 2. Gather the image K/V over the orthogonal group, then let the + # inherited AllGather-KV path slice the attention metadata, prepend + # the joint tensors, and build the local-Q/global-KV view. + self._assert_equal_region_lengths(key.shape[1], key.device) + query, key, value, attn_metadata, _ = super().pre_attention(query, key, value, attn_metadata) + + # 3. Own all 2D key-mask handling against the final key layout: merge + # joint/image masks (filling missing sides), or leave 4D masks to + # the inherited query-range slicing. + attn_metadata = self._merge_joint_attn_mask(attn_metadata, key) + + # 4. The reverse transform is entirely Ulysses': it splits the joint + # part back out, undoes the image all-to-all, and head-gathers the + # joint output over the Ulysses group. + return query, key, value, attn_metadata, ctx + + def post_attention( + self, + attn_output: torch.Tensor, + ctx: ParallelAttentionContext | None, + ) -> torch.Tensor: + return self._ulysses.post_attention(attn_output, ctx) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 39d5449629b..53328e1a6b6 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -267,8 +267,7 @@ class DiffusionParallelConfig: sequence_parallel_size: int | None = None """Number of sequence parallel groups. - sequence_parallel_size = ulysses_degree * ring_degree, or allgather_degree - when AllGather-KV is enabled.""" + sequence_parallel_size = ulysses_degree * ring_degree * allgather_degree.""" ulysses_degree: int = 1 """Number of GPUs used for ulysses sequence parallelism.""" @@ -277,7 +276,8 @@ class DiffusionParallelConfig: """Number of GPUs used for ring sequence parallelism.""" allgather_degree: int = 1 - """Number of GPUs used for AllGather-KV sequence parallelism (causal=False only).""" + """Number of GPUs used for AllGather-KV sequence parallelism (causal=False only). + Composable with ``ulysses_degree``; not with ``ring_degree > 1``.""" ulysses_mode: str = "strict" """Ulysses sequence-parallel mode. @@ -367,15 +367,15 @@ def _validate_parallel_config(self) -> Self: self.data_parallel_size not in (None, 1) or self.pipeline_parallel_size != 1 or self.cfg_parallel_size != 1 ): raise ValueError("VAE batch parallel decode requires DP, PP, and CFG parallel sizes to be 1") - if self.allgather_degree > 1: - assert self.ulysses_degree == 1 and self.ring_degree == 1, ( - "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. " + if self.allgather_degree > 1 and self.ring_degree > 1: + raise ValueError( + "AllGather-KV (allgather_degree>1) cannot be composed with Ring (ring_degree>1). " + "The supported two-dimensional topology is Ulysses x AllGather-KV " + "(ulysses_degree > 1 with ring_degree == 1). " f"Got ulysses_degree={self.ulysses_degree}, ring_degree={self.ring_degree}, " f"allgather_degree={self.allgather_degree}." ) - expected_sp_size = ( - self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree - ) + expected_sp_size = self.ulysses_degree * self.ring_degree * self.allgather_degree assert self.sequence_parallel_size == expected_sp_size, ( f"Sequence parallel size must be {expected_sp_size}, but got {self.sequence_parallel_size}" ) @@ -391,9 +391,7 @@ def _validate_parallel_config(self) -> Self: def __post_init__(self) -> None: if self.sequence_parallel_size is None: - self.sequence_parallel_size = ( - self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree - ) + self.sequence_parallel_size = self.ulysses_degree * self.ring_degree * self.allgather_degree # Until the runtime WORLD size is known, an omitted DP dimension means # one replica. OmniDiffusionConfig resolves it against num_gpus below. diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 858a2e398c9..8cdca01d49d 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -221,9 +221,9 @@ def __init__( def get_mask(self, order: str, token: str): ordered_token = order.split("-") - token = token.split("-") + tokens = token.split("-") mask = [False] * len(ordered_token) - for t in token: + for t in tokens: mask[ordered_token.index(t)] = True return mask @@ -523,6 +523,87 @@ def init_vllm_model_parallel_group( ) +def build_ulysses_allgather_rank_groups( + sp_group_ranks: list[list[int]], + sp_ulysses_degree: int, + sp_allgather_degree: int, +) -> tuple[list[list[int]], list[list[int]]]: + """Rank lists of every Ulysses / AllGather-KV group in the composed topology. + + Pure function (no ``torch.distributed``). Position ``p`` inside one SP + group maps to ``ulysses_rank = p % U`` and ``allgather_rank = p // U``:: + + SP group [0 1 2 3] (U=2, A=2): Ulysses [0 1] [2 3] AllGather [0 2] [1 3] + + This is the layout invariant of the composed topology, chosen to match the + flat contiguous sharding of ``sp_shard`` (position ``p`` owns chunk ``p``): + a contiguous Ulysses group jointly holds one contiguous region, so its + all-to-all reproduces that region in order, and the stride-U AllGather + group holds regions in ``allgather_rank`` order, so the K/V all-gather + rebuilds the global sequence. Any other layout permutes the sequence. + + Returns: + ``(ulysses_groups, allgather_groups)`` in a deterministic order, so + every rank issues the same ``new_group`` sequence. + """ + ulysses_groups: list[list[int]] = [] + allgather_groups: list[list[int]] = [] + for group_ranks in sp_group_ranks: + if len(group_ranks) != sp_ulysses_degree * sp_allgather_degree: + raise ValueError( + f"Invalid sp_group_ranks entry: expected size " + f"{sp_ulysses_degree * sp_allgather_degree}, got {len(group_ranks)}." + ) + # Ulysses groups: contiguous regions of U ranks. One region per + # AllGather rank, so region index == allgather_rank. + for region in range(sp_allgather_degree): + ulysses_groups.append(list(group_ranks[region * sp_ulysses_degree : (region + 1) * sp_ulysses_degree])) + # AllGather-KV groups: stride-U, so index inside the group == allgather_rank. + for ulysses_rank in range(sp_ulysses_degree): + allgather_groups.append(list(group_ranks[ulysses_rank::sp_ulysses_degree])) + return ulysses_groups, allgather_groups + + +def _set_ulysses_allgather_pg( + sp_ulysses_degree: int, + sp_allgather_degree: int, + rank: int, + world_size: int, + sp_group_ranks: list[list[int]] | None, +) -> tuple[torch.distributed.ProcessGroup, torch.distributed.ProcessGroup, torch.distributed.ProcessGroup]: + """Build the orthogonal process groups of the Ulysses x AllGather-KV topology. + + See ``build_ulysses_allgather_rank_groups`` for the rank-layout invariant. + ``ring_degree == 1`` is required, so every rank's Ring group is a singleton. + """ + sp_size = sp_ulysses_degree * sp_allgather_degree + if sp_group_ranks is None: + sp_group_ranks = [list(range(offset, offset + sp_size)) for offset in range(0, world_size, sp_size)] + if len(sp_group_ranks) * sp_size != world_size or any(len(ranks) != sp_size for ranks in sp_group_ranks): + raise ValueError(f"Invalid sp_group_ranks: expected {world_size // sp_size} groups of size {sp_size}.") + + ulysses_pg = ring_pg = allgather_pg = None + ulysses_groups, allgather_groups = build_ulysses_allgather_rank_groups( + sp_group_ranks, + sp_ulysses_degree, + sp_allgather_degree, + ) + for ranks in ulysses_groups: + group = torch.distributed.new_group(ranks) + if rank in ranks: + ulysses_pg = group + for ranks in allgather_groups: + group = torch.distributed.new_group(ranks) + if rank in ranks: + allgather_pg = group + for singleton_rank in range(world_size): + group = torch.distributed.new_group([singleton_rank]) + if rank == singleton_rank: + ring_pg = group + assert ulysses_pg is not None and ring_pg is not None and allgather_pg is not None + return ulysses_pg, ring_pg, allgather_pg + + # adapted from https://github.com/feifeibear/long-context-attention/blob/main/yunchang/globals.py def set_seq_parallel_pg( sp_ulysses_degree: int, @@ -536,7 +617,8 @@ def set_seq_parallel_pg( """ Initialize Ulysses, Ring, and AllGather-KV process groups. - AllGather-KV is mutually exclusive with Ulysses and Ring. + AllGather-KV may be composed with Ulysses (the orthogonal + Ulysses x AllGather-KV topology), but not with Ring. Args: sp_ulysses_degree: Size of each Ulysses subgroup. @@ -546,19 +628,22 @@ def set_seq_parallel_pg( use_ulysses_low: If True, Ulysses groups are contiguous chunks and Ring groups are strided within each SP group. If False, the opposite. sp_group_ranks: Optional explicit SP groups. Each entry must be a list - of length sp_ulysses_degree * sp_ring_degree. When provided, groups - are built from these ranks instead of auto-generated contiguous - ranges. + of length sp_ulysses_degree * sp_ring_degree * sp_allgather_degree. + When provided, groups are built from these ranks instead of + auto-generated contiguous ranges. Returns: ulysses_pg (torch.distributed.ProcessGroup): The Ulysses process group for this rank. ring_pg (torch.distributed.ProcessGroup): The Ring process group for this rank. + allgather_pg (torch.distributed.ProcessGroup): The AllGather-KV process + group for this rank. Raises: - ValueError: If sp_group_ranks length does not match world_size or any - entry has the wrong size. + ValueError: If sp_group_ranks length does not match world_size, any + entry has the wrong size, or an unsupported degree combination is + requested. AssertionError: If world_size is not divisible by sp_size. Behavior: @@ -569,8 +654,26 @@ def set_seq_parallel_pg( slice using offsets of size sp_size. """ if sp_allgather_degree > 1: - if sp_ulysses_degree > 1 or sp_ring_degree > 1: - raise ValueError("AllGather-KV is mutually exclusive with Ulysses and Ring") + if sp_ring_degree > 1: + raise ValueError( + "AllGather-KV cannot be composed with Ring: the supported two-dimensional " + "topology is Ulysses x AllGather-KV (ulysses_degree > 1 with ring_degree == 1). " + f"Got ulysses_degree={sp_ulysses_degree}, ring_degree={sp_ring_degree}, " + f"allgather_degree={sp_allgather_degree}." + ) + if sp_ulysses_degree > 1: + logger.info( + "Building orthogonal Ulysses x AllGather-KV process groups " + f"(ulysses={sp_ulysses_degree}, allgather={sp_allgather_degree})." + ) + return _set_ulysses_allgather_pg( + sp_ulysses_degree=sp_ulysses_degree, + sp_allgather_degree=sp_allgather_degree, + rank=rank, + world_size=world_size, + sp_group_ranks=sp_group_ranks, + ) + sp_size = sp_allgather_degree if sp_group_ranks is None: sp_group_ranks = [list(range(offset, offset + sp_size)) for offset in range(0, world_size, sp_size)] @@ -726,12 +829,13 @@ def _initialize_model_parallel( data_parallel_size: number of data parallelism groups. cfg_parallel_size: number of GPUs used for Classifier Free Guidance (CFG) parallelism. sequence_parallel_size: number of GPUs used for sequence parallelism. - Uses allgather_degree when AllGather-KV is enabled, otherwise - ulysses_degree * ring_degree. + Equals ``ulysses_degree * ring_degree * allgather_degree``. ulysses_degree: number of GPUs used for ulysses sequence parallelism. ring_degree: number of GPUs used for ring sequence parallelism. allgather_degree: number of GPUs used for AllGather-KV sequence parallelism - (causal=False only). Mutually exclusive with ulysses/ring in v1. + (causal=False only). May be composed with ulysses_degree as the + orthogonal Ulysses x AllGather-KV topology; cannot be combined with + ring_degree > 1. tensor_parallel_size: number of GPUs used for tensor parallelism. pipeline_parallel_size: number of GPUs used for pipeline parallelism. fully_shard_degree: number of GPUs used for the HSDP shard dimension. @@ -771,15 +875,16 @@ def _initialize_model_parallel( world_size: int = torch.distributed.get_world_size() backend = backend or torch.distributed.get_backend(get_world_group().device_group) - if allgather_degree > 1: - if ulysses_degree != 1 or ring_degree != 1: - raise ValueError( - "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. " - f"Got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, " - f"allgather_degree={allgather_degree}." - ) + if allgather_degree > 1 and ring_degree > 1: + raise ValueError( + "AllGather-KV (allgather_degree>1) cannot be composed with Ring (ring_degree>1). " + "The supported two-dimensional topology is Ulysses x AllGather-KV " + "(ulysses_degree > 1 with ring_degree == 1). " + f"Got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, " + f"allgather_degree={allgather_degree}." + ) - expected_sequence_parallel_size = allgather_degree if allgather_degree > 1 else ring_degree * ulysses_degree + expected_sequence_parallel_size = ring_degree * ulysses_degree * allgather_degree if sequence_parallel_size is None: sequence_parallel_size = expected_sequence_parallel_size logger.info("sequence_parallel_size is not provided, using %d", sequence_parallel_size) @@ -856,7 +961,7 @@ def get_rank_groups(token: str) -> list[list[int]]: ) global _PP assert _PP is None, "pipeline model parallel group is already initialized" - _PP = init_model_parallel_group( + _PP = init_model_parallel_group( # type: ignore[assignment] group_ranks=get_rank_groups("pp"), local_rank=get_world_group().local_rank, backend=backend, @@ -874,7 +979,7 @@ def get_rank_groups(token: str) -> list[list[int]]: world_size=world_size, sp_group_ranks=sp_group_ranks, ) - _SP = init_model_parallel_group( + _SP = init_model_parallel_group( # type: ignore[assignment] group_ranks=sp_group_ranks, local_rank=get_world_group().local_rank, backend=backend, diff --git a/vllm_omni/diffusion/distributed/sp_plan.py b/vllm_omni/diffusion/distributed/sp_plan.py index e4f97e85219..f6ea546b14a 100644 --- a/vllm_omni/diffusion/distributed/sp_plan.py +++ b/vllm_omni/diffusion/distributed/sp_plan.py @@ -68,12 +68,17 @@ class SequenceParallelConfig: ring_degree: Number of devices for Ring attention. Sequence is split across devices, with K/V passed in a ring topology. Best for long sequences with limited memory/bandwidth. + allgather_degree: Number of devices for AllGather-KV attention. Sequence is + split across devices and K/V are gathered in an orthogonal group, so + every rank runs local-Q/global-KV attention. convert_to_fp32: Whether to convert output and LSE to float32 for numerical stability in ring attention. Note: - ulysses_degree * ring_degree = sequence_parallel_size - vLLM-Omni supports hybrid Ulysses-Ring attention (both > 1). + ulysses_degree * ring_degree * allgather_degree = sequence_parallel_size + vLLM-Omni supports hybrid Ulysses-Ring attention (both > 1) and the + orthogonal Ulysses x AllGather-KV topology (ulysses_degree and + allgather_degree > 1 with ring_degree == 1). """ ulysses_degree: int = 1 @@ -90,15 +95,18 @@ def __post_init__(self) -> None: if self.ulysses_degree < 1 or self.ring_degree < 1 or self.allgather_degree < 1: raise ValueError("SP degrees must be >= 1.") - if self.allgather_degree > 1 and (self.ulysses_degree > 1 or self.ring_degree > 1): - raise ValueError("AllGather-KV is mutually exclusive with Ulysses and Ring.") + if self.allgather_degree > 1 and self.ring_degree > 1: + raise ValueError( + "AllGather-KV cannot be composed with Ring; the supported two-dimensional " + "topology is Ulysses x AllGather-KV (ulysses_degree > 1 with ring_degree == 1)." + ) if self.ulysses_degree == self.ring_degree == self.allgather_degree == 1: raise ValueError("At least one SP degree must be > 1.") @property def sequence_parallel_size(self) -> int: """Total sequence parallel world size.""" - return self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree + return self.ulysses_degree * self.ring_degree * self.allgather_degree def get_world_size(self) -> int: """Get the sequence parallel world size from parallel state. diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 27f8d78892a..3ed45258491 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -621,28 +621,19 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - if plan is None: continue - # AllGather-KV reuses the Ulysses sequence-sharding hooks. - allgather_degree = getattr(od_config.parallel_config, "allgather_degree", 1) - if allgather_degree > 1: - sp_config = SequenceParallelConfig( - allgather_degree=allgather_degree, - ) - mode = "allgather_kv" - else: - sp_config = SequenceParallelConfig( - ulysses_degree=od_config.parallel_config.ulysses_degree, - ring_degree=od_config.parallel_config.ring_degree, - ) - # Apply hooks according to the plan - mode = ( - "hybrid" - if sp_config.ulysses_degree > 1 and sp_config.ring_degree > 1 - else ("ulysses" if sp_config.ulysses_degree > 1 else "ring") - ) + # The SP hooks shard by `sequence_parallel_size`, so carry the full + # degree triple; AllGather-KV and the composed topology reuse them as-is. + parallel_config = od_config.parallel_config + sp_config = SequenceParallelConfig( + ulysses_degree=parallel_config.ulysses_degree, + ring_degree=parallel_config.ring_degree, + allgather_degree=parallel_config.allgather_degree, + ) logger.info( f"Applying sequence parallelism to {transformer.__class__.__name__} ({attr}) " - f"(sp_size={sp_size}, mode={mode}, ulysses={sp_config.ulysses_degree}, ring={sp_config.ring_degree})" + f"(sp_size={sp_size}, ulysses={sp_config.ulysses_degree}, ring={sp_config.ring_degree}, " + f"allgather={sp_config.allgather_degree})" ) apply_sequence_parallel(transformer, sp_config, plan) applied_count += 1