[Feature][Diffusion] Add BOOGU Image sequence parallel support (#5718)

Signed-off-by: Xenoryn <timberh2000@outlook.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/attention/test_ulysses_uaa.py

@@ -5,13 +5,16 @@ from __future__ import annotations
 
 import os
 import tempfile
+from contextlib import AbstractContextManager
 from types import SimpleNamespace
 
 import numpy as np
 import pytest
 import torch
 
+from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
 from vllm_omni.diffusion.config import set_current_diffusion_config
@@ -578,6 +581,7 @@ def _run_attention_case(
     num_kv_heads: int | None = None,
     has_joint: bool = False,
     joint_strategy: str = "front",
+    mask_layout: str | None = None,
 ) -> None:
     device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
     current_omni_platform.set_device(device)
@@ -620,6 +624,7 @@ def _run_attention_case(
             q_full = torch.from_numpy(payload["q"]).to(device=device)
             k_full = torch.from_numpy(payload["k"]).to(device=device)
             v_full = torch.from_numpy(payload["v"]).to(device=device)
+            mask_full = torch.from_numpy(payload["mask"]).to(device=device) if "mask" in payload else None
             if has_joint:
                 joint_q = torch.from_numpy(payload["joint_q"]).to(device=device)
                 joint_k = torch.from_numpy(payload["joint_k"]).to(device=device)
@@ -640,6 +645,7 @@ def _run_attention_case(
                     v = torch.cat([v_full, joint_v], dim=1)
             else:
                 q, k, v = q_full, k_full, v_full
+            mask = mask_full
         else:
             if split_sizes is None:
                 # NOTE: torch.chunk may return fewer than `world_size` chunks for some
@@ -659,14 +665,18 @@ def _run_attention_case(
                 q = torch.split(q_full, split_sizes, dim=1)[local_rank].contiguous()
                 k = torch.split(k_full, split_sizes, dim=1)[local_rank].contiguous()
                 v = torch.split(v_full, split_sizes, dim=1)[local_rank].contiguous()
+            if mask_full is None:
+                mask = None
+            elif mask_layout == "global":
+                mask = mask_full
+            else:
+                raise ValueError("mask_layout must be 'global' when a mask is provided.")
             # The Attention layer only enables SP communication when ForwardContext.sp_active is True.
             # In production this is managed by SequenceParallelSplitHook/GatherHook, but here we
             # shard manually for testing.
             get_forward_context()._sp_shard_depth = 1
 
             if has_joint:
-                from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
-
                 attn_metadata = AttentionMetadata(
                     joint_query=joint_q,
                     joint_key=joint_k,
@@ -682,6 +692,8 @@ def _run_attention_case(
             raise ValueError(f"Invalid sdp_kernel_mode: {sdp_kernel_mode!r}")
 
         with sdp_ctx:
+            if mask is not None:
+                attn_metadata = AttentionMetadata(attn_mask=mask)
             out_local = attn(q, k, v, attn_metadata=attn_metadata).contiguous()
 
         if world_size == 1:
@@ -908,7 +920,7 @@ def _run_contract_skip_case(
         ctx = get_forward_context()
         ctx._sp_shard_depth = 1
 
-        def make_sdp_ctx() -> object:
+        def make_sdp_ctx() -> AbstractContextManager[None]:
             # Fresh instance per phase: the @contextmanager instance releases
             # its state on first exit and cannot be re-entered.
             return torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True)
@@ -1011,6 +1023,105 @@ def test_ulysses_uaa_contract_skips_length_gather_matches_slow_path(num_heads: i
             pass
 
 
+@pytest.mark.core_model
+@pytest.mark.parametrize("mask_layout", ["global"])
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+def test_ulysses_uaa_2d_mask_layout_matches_baseline(
+    mask_layout: str,
+) -> None:
+    world_size = 2
+    batch_size = 2
+    seq_len = 5
+    split_sizes = [2, 3]
+    num_heads = 3
+    head_size = 8
+
+    base_init_method = get_distributed_init_method()
+    sp_init_method = get_distributed_init_method()
+
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in:
+        input_file = f_in.name
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".npy") as f_base:
+        baseline_file = f_base.name
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".npy") as f_sp:
+        sp_file = f_sp.name
+
+    try:
+        torch.manual_seed(0)
+        q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
+        k = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
+        v = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
+        mask = torch.tensor(
+            [
+                [True, True, True, True, False],
+                [True, True, True, False, False],
+            ],
+            dtype=torch.bool,
+        )
+        np.savez(
+            input_file,
+            q=q.numpy(),
+            k=k.numpy(),
+            v=v.numpy(),
+            mask=mask.numpy(),
+        )
+
+        # Baseline (no SP), mask is used verbatim.
+        torch.multiprocessing.spawn(
+            _run_attention_case,
+            args=(
+                1,
+                base_init_method,
+                input_file,
+                baseline_file,
+                num_heads,
+                head_size,
+                1,
+                "strict",
+                1,
+                None,
+                "math",
+                None,
+                False,
+                "front",
+                "global",
+            ),
+            nprocs=1,
+        )
+        # SP (Ulysses-P with UAA), mask is shared globally across ranks.
+        torch.multiprocessing.spawn(
+            _run_attention_case,
+            args=(
+                world_size,
+                sp_init_method,
+                input_file,
+                sp_file,
+                num_heads,
+                head_size,
+                world_size,
+                "advanced_uaa",
+                1,
+                split_sizes,
+                "math",
+                None,
+                False,
+                "front",
+                mask_layout,
+            ),
+            nprocs=world_size,
+        )
+
+        baseline = torch.from_numpy(np.load(baseline_file, allow_pickle=False))
+        sp = torch.from_numpy(np.load(sp_file, allow_pickle=False))
+        torch.testing.assert_close(sp, baseline, atol=1e-5, rtol=1e-5)
+    finally:
+        for path in (input_file, baseline_file, sp_file):
+            try:
+                os.remove(path)
+            except OSError:
+                pass
+
+
 @pytest.mark.parametrize(
     "num_heads,num_kv_heads,joint_len",
     [

tests/diffusion/distributed/test_boogu_image_sp.py

@@ -0,0 +1,317 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import os
+import socket
+import tempfile
+
+import numpy as np
+import pytest
+import torch
+from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import (
+    DiffusionParallelConfig,
+    OmniDiffusionConfig,
+    TransformerConfig,
+)
+from vllm_omni.diffusion.distributed.parallel_state import (
+    destroy_distributed_env,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+from vllm_omni.diffusion.forward_context import (
+    get_forward_context,
+    set_forward_context,
+)
+from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.diffusion,
+    pytest.mark.parallel,
+]
+
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return int(sock.getsockname()[1])
+
+
+def _set_dist_env(rank: int, world_size: int, port: int) -> None:
+    os.environ.update(
+        {
+            "RANK": str(rank),
+            "LOCAL_RANK": str(rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "127.0.0.1",
+            "MASTER_PORT": str(port),
+            "DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA",
+        }
+    )
+
+
+def _transformer_config() -> TransformerConfig:
+    return TransformerConfig.from_dict(
+        {
+            "patch_size": 2,
+            "in_channels": 4,
+            "hidden_size": 64,
+            "num_layers": 2,
+            "num_double_stream_layers": 1,
+            "num_refiner_layers": 1,
+            "num_attention_heads": 4,
+            "num_kv_heads": 1,
+            "multiple_of": 32,
+            "norm_eps": 1e-5,
+            "axes_dim_rope": [8, 4, 4],
+            "axes_lens": [64, 32, 32],
+            "instruction_feature_configs": {
+                "instruction_feat_dim": 32,
+                "reduce_type": "mean",
+                "num_instruction_feature_layers": 1,
+            },
+            "prompt_tuning_configs": {"use_prompt_tuning": False},
+            "timestep_scale": 1.0,
+        }
+    )
+
+
+def _od_config(world_size: int) -> OmniDiffusionConfig:
+    return OmniDiffusionConfig(
+        model="boogu-tiny",
+        dtype=torch.float32,
+        tf_model_config=_transformer_config(),
+        parallel_config=DiffusionParallelConfig(
+            sequence_parallel_size=world_size,
+            ulysses_degree=world_size,
+            ulysses_mode="advanced_uaa",
+        ),
+    )
+
+
+def _checkpoint_name(name: str) -> str:
+    if ".to_out." in name:
+        name = name.replace(".to_out.", ".to_out.0.")
+    for projection in (
+        "img_to_q",
+        "img_to_k",
+        "img_to_v",
+        "instruct_to_q",
+        "instruct_to_k",
+        "instruct_to_v",
+        "instruct_out",
+        "img_out",
+    ):
+        token = f".img_instruct_attn.{projection}."
+        if token in name:
+            return name.replace(
+                token,
+                f".img_instruct_attn.processor.{projection}.",
+            )
+    return name
+
+
+def _random_checkpoint(model: torch.nn.Module) -> list[tuple[str, torch.Tensor]]:
+    generator = torch.Generator(device=next(model.parameters()).device).manual_seed(2026)
+    return [
+        (
+            _checkpoint_name(name),
+            torch.empty_like(param).uniform_(
+                -0.02,
+                0.02,
+                generator=generator,
+            ),
+        )
+        for name, param in model.named_parameters()
+    ]
+
+
+def _worker(
+    rank: int,
+    world_size: int,
+    port: int,
+    output_file: str,
+) -> None:
+    device = torch.device(f"{current_omni_platform.device_type}:{rank}")
+    current_omni_platform.set_device(device)
+    _set_dist_env(rank, world_size, port)
+    init_distributed_environment(world_size=world_size, rank=rank)
+    initialize_model_parallel(
+        data_parallel_size=1,
+        cfg_parallel_size=1,
+        sequence_parallel_size=world_size,
+        ulysses_degree=world_size,
+        ring_degree=1,
+        tensor_parallel_size=1,
+        pipeline_parallel_size=1,
+    )
+
+    try:
+        from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+            BooguImageDoubleStreamRotaryPosEmbed,
+            BooguImageTransformer2DModel,
+        )
+
+        od_config = _od_config(world_size)
+        with (
+            set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device=str(device)))),
+            set_forward_context(omni_diffusion_config=od_config),
+            set_current_diffusion_config(od_config),
+        ):
+            source = BooguImageTransformer2DModel(od_config).to(device)
+            checkpoint = _random_checkpoint(source)
+            del source
+
+            model = BooguImageTransformer2DModel(od_config).to(device)
+            assert model.load_weights(checkpoint) == set(dict(model.named_parameters()))
+            model.eval()
+
+            if world_size > 1:
+                apply_sequence_parallel(
+                    model,
+                    SequenceParallelConfig(ulysses_degree=world_size),
+                    model._sp_plan,
+                )
+                get_forward_context().sp_plan_hooks_applied = True
+
+            generator = torch.Generator(device=device).manual_seed(17)
+            latents = torch.randn(
+                1,
+                4,
+                8,
+                8,
+                device=device,
+                generator=generator,
+            )
+            timestep = torch.tensor([0.5], device=device)
+            instruction = torch.randn(
+                1,
+                7,
+                32,
+                device=device,
+                generator=generator,
+            )
+            instruction_mask = torch.ones(
+                1,
+                7,
+                dtype=torch.bool,
+                device=device,
+            )
+            ref_image = [
+                [
+                    torch.randn(
+                        4,
+                        6,
+                        10,
+                        device=device,
+                        generator=generator,
+                    )
+                ]
+            ]
+            # Two differently sized references: their packed segments straddle
+            # the SP rank boundary, which is what the per-rank reference-refiner
+            # mask has to reproduce.
+            ref_images_multi = [
+                [
+                    torch.randn(4, 6, 10, device=device, generator=generator),
+                    torch.randn(4, 4, 8, device=device, generator=generator),
+                ]
+            ]
+            freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(
+                model.axes_dim_rope,
+                model.axes_lens,
+                theta=10000,
+            )
+
+            with torch.inference_mode():
+                t2i = model(
+                    latents,
+                    timestep,
+                    instruction,
+                    freqs_real,
+                    instruction_mask,
+                )
+                edit = model(
+                    latents,
+                    timestep,
+                    instruction,
+                    freqs_real,
+                    instruction_mask,
+                    ref_image_hidden_states=ref_image,
+                )
+                edit_multi = model(
+                    latents,
+                    timestep,
+                    instruction,
+                    freqs_real,
+                    instruction_mask,
+                    ref_image_hidden_states=ref_images_multi,
+                )
+
+            if rank == 0:
+                np.savez(
+                    output_file,
+                    t2i=t2i.cpu().numpy(),
+                    edit=edit.cpu().numpy(),
+                    edit_multi=edit_multi.cpu().numpy(),
+                )
+    finally:
+        destroy_distributed_env()
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+def test_boogu_sp_matches_single_gpu() -> None:
+    if not current_omni_platform.is_cuda() or current_omni_platform.get_device_count() < 2:
+        pytest.skip("BOOGU SP test requires two CUDA GPUs.")
+
+    output_files = []
+    try:
+        for _ in range(2):
+            with tempfile.NamedTemporaryFile(
+                delete=False,
+                suffix=".npz",
+            ) as handle:
+                output_files.append(handle.name)
+
+        torch.multiprocessing.spawn(
+            _worker,
+            args=(1, _find_free_port(), output_files[0]),
+            nprocs=1,
+        )
+        torch.multiprocessing.spawn(
+            _worker,
+            args=(2, _find_free_port(), output_files[1]),
+            nprocs=2,
+        )
+
+        with (
+            np.load(output_files[0], allow_pickle=False) as baseline,
+            np.load(output_files[1], allow_pickle=False) as sp_output,
+        ):
+            for key in ("t2i", "edit", "edit_multi"):
+                # fp32 SDPA: the observed SP1-vs-SP2 gap is ~3e-8, so 2e-4
+                # would also pass with the SP masking logic entirely removed
+                # (that mutation measures ~5e-7). Mask correctness is pinned
+                # structurally in tests/diffusion/models/boogu_image/
+                # test_sp_layout.py; this bound just stays close to the real
+                # reduction-order noise instead of hiding logic errors.
+                np.testing.assert_allclose(
+                    sp_output[key],
+                    baseline[key],
+                    rtol=1e-5,
+                    atol=1e-5,
+                )
+    finally:
+        for path in output_files:
+            try:
+                os.remove(path)
+            except OSError:
+                pass

tests/diffusion/distributed/test_sp_plan_hooks.py

@@ -115,6 +115,47 @@ class TestSequenceParallelPlanValidation:
         with pytest.raises(ValueError, match="split_output=True"):
             validate_sp_plan(plan)
 
+    def test_keyed_metadata_requires_auto_pad(self):
+        plan = {
+            "prepare": {
+                0: SequenceParallelInput(
+                    split_dim=1,
+                    split_output=True,
+                    shard_group="image",
+                ),
+            },
+            "output": SequenceParallelOutput(
+                gather_dim=1,
+                shard_group="image",
+            ),
+        }
+
+        with pytest.raises(ValueError, match="requires auto_pad=True"):
+            validate_sp_plan(plan)
+
+    def test_shard_group_is_scoped_to_one_split_boundary(self):
+        plan = {
+            "prepare_image": {
+                0: SequenceParallelInput(
+                    split_dim=1,
+                    split_output=True,
+                    auto_pad=True,
+                    shard_group="image",
+                ),
+            },
+            "prepare_rope": {
+                0: SequenceParallelInput(
+                    split_dim=1,
+                    split_output=True,
+                    auto_pad=True,
+                    shard_group="image",
+                ),
+            },
+        }
+
+        with pytest.raises(ValueError, match="multiple split boundaries"):
+            validate_sp_plan(plan)
+
 
 @pytest.mark.cpu
 class TestGetSpPlanFromModel:
@@ -1006,6 +1047,8 @@ class TestDimensionValidation:
                 return "NO_MASK_BACKEND"
 
         monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 0)
+        monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
         monkeypatch.setattr(selector, "get_attn_backend_for_capability", lambda **kwargs: NoMaskBackend)
 
         hook = SequenceParallelSplitHook(
@@ -1015,7 +1058,7 @@ class TestDimensionValidation:
         tensor = torch.randn(1, 3, 8)
 
         with pytest.raises(ValueError, match="does not support attention_mask"):
-            hook._shard_with_auto_pad(tensor, dim=1)
+            hook._shard_with_auto_pad(tensor, dim=1, shard_group=None)
 
     def test_auto_pad_allows_explicit_cudnn(self, monkeypatch):
         from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
@@ -1040,7 +1083,7 @@ class TestDimensionValidation:
         )
 
         with set_forward_context(omni_diffusion_config=od_config):
-            result = hook._shard_with_auto_pad(tensor, dim=1)
+            result = hook._shard_with_auto_pad(tensor, dim=1, shard_group=None)
 
         # seq_len=3 padded to 4, then sharded by world_size=2 → rank 0 gets 2.
         assert result.shape == (1, 2, 8)
@@ -1070,7 +1113,7 @@ class TestDimensionValidation:
 
         with set_forward_context(omni_diffusion_config=od_config):
             with pytest.raises(ValueError, match="does not support attention_mask"):
-                hook._shard_with_auto_pad(tensor, dim=1)
+                hook._shard_with_auto_pad(tensor, dim=1, shard_group=None)
 
 
 @pytest.mark.cpu
@@ -1227,6 +1270,158 @@ class TestEqualPadContract:
             assert not ctx.sp_rank_local_seq_lens_equal
 
 
+@pytest.mark.cpu
+class TestKeyedShardMetadata:
+    """Per-shard_group padding metadata, for models with several sequences.
+
+    A model that splits image, reference and context tokens at one boundary
+    needs each one's pre-padding global length to rebuild masks; the singleton
+    sp_original_seq_len can only describe one of them.
+    """
+
+    def test_auto_pad_records_the_global_length_under_its_key(self, monkeypatch):
+        from vllm_omni.diffusion.attention import selector
+        from vllm_omni.diffusion.distributed import parallel_state
+        from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+        from vllm_omni.diffusion.forward_context import (
+            get_forward_context,
+            get_sp_shard_original_seq_len,
+            set_forward_context,
+        )
+        from vllm_omni.diffusion.hooks.sequence_parallel import (
+            SequenceParallelSplitHook,
+        )
+
+        class MaskBackend:
+            supports_attention_mask = True
+
+            @staticmethod
+            def get_name():
+                return "test"
+
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 1)
+        monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+        monkeypatch.setattr(selector, "get_attn_backend_for_role", lambda **_: (MaskBackend(), None))
+
+        hook = SequenceParallelSplitHook(
+            {
+                0: SequenceParallelInput(
+                    split_dim=1,
+                    expected_dims=3,
+                    split_output=True,
+                    auto_pad=True,
+                    shard_group="image",
+                )
+            },
+            SequenceParallelConfig(ulysses_degree=2),
+        )
+        hook.initialize_hook(nn.Identity())
+
+        with set_forward_context():
+            hook.pre_forward(nn.Identity())
+            shard = hook.post_forward(nn.Identity(), torch.arange(5).reshape(1, 5, 1))
+            ctx = get_forward_context()
+
+            assert shard.shape[1] == 3
+            # The *pre*-padding length, which is what a mask has to clip to.
+            assert ctx.sp_shard_metadata["image"] == 5
+            assert get_sp_shard_original_seq_len("image") == 5
+            assert get_sp_shard_original_seq_len("context") is None
+
+    def test_gather_trims_to_its_key_and_leaves_other_groups(self, monkeypatch):
+        from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+        from vllm_omni.diffusion.forward_context import (
+            get_forward_context,
+            set_forward_context,
+        )
+        from vllm_omni.diffusion.hooks import sequence_parallel
+        from vllm_omni.diffusion.hooks.sequence_parallel import (
+            SequenceParallelGatherHook,
+        )
+
+        monkeypatch.setattr(
+            sequence_parallel,
+            "sp_gather",
+            lambda tensor, dim, validate: torch.cat([tensor, tensor], dim=dim),
+        )
+        hook = SequenceParallelGatherHook(
+            SequenceParallelOutput(gather_dim=1, expected_dims=3, shard_group="image"),
+            SequenceParallelConfig(ulysses_degree=2),
+        )
+
+        with set_forward_context():
+            ctx = get_forward_context()
+            ctx.sp_shard_metadata["image"] = 5
+            ctx.sp_shard_metadata["context"] = 7
+            ctx._sp_shard_depth = 1
+            output = hook.post_forward(nn.Identity(), torch.zeros(1, 3, 4))
+
+            # 2x3 gathered, trimmed back to "image"'s recorded 5.
+            assert output.shape == (1, 5, 4)
+            # A concurrent group is untouched.
+            assert ctx.sp_shard_metadata == {"context": 7}
+
+    def test_gather_without_registered_metadata_is_an_error(self, monkeypatch):
+        from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+        from vllm_omni.diffusion.forward_context import set_forward_context
+        from vllm_omni.diffusion.hooks import sequence_parallel
+        from vllm_omni.diffusion.hooks.sequence_parallel import (
+            SequenceParallelGatherHook,
+        )
+
+        monkeypatch.setattr(
+            sequence_parallel,
+            "sp_gather",
+            lambda tensor, dim, validate: torch.cat([tensor, tensor], dim=dim),
+        )
+        hook = SequenceParallelGatherHook(
+            SequenceParallelOutput(gather_dim=1, expected_dims=3, shard_group="image"),
+            SequenceParallelConfig(ulysses_degree=2),
+        )
+
+        with set_forward_context(), pytest.raises(RuntimeError, match="shard_group='image'"):
+            hook.post_forward(nn.Identity(), torch.zeros(1, 3, 4))
+
+    def test_conflicting_lengths_in_one_group_are_rejected(self, monkeypatch):
+        from vllm_omni.diffusion.attention import selector
+        from vllm_omni.diffusion.distributed import parallel_state
+        from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+        from vllm_omni.diffusion.forward_context import set_forward_context
+        from vllm_omni.diffusion.hooks.sequence_parallel import (
+            SequenceParallelSplitHook,
+        )
+
+        class MaskBackend:
+            supports_attention_mask = True
+
+            @staticmethod
+            def get_name():
+                return "test"
+
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 1)
+        monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+        monkeypatch.setattr(selector, "get_attn_backend_for_role", lambda **_: (MaskBackend(), None))
+
+        spm = SequenceParallelInput(
+            split_dim=1,
+            expected_dims=3,
+            split_output=True,
+            auto_pad=True,
+            shard_group="image",
+        )
+        hook = SequenceParallelSplitHook({0: spm, 1: spm}, SequenceParallelConfig(ulysses_degree=2))
+        hook.initialize_hook(nn.Identity())
+
+        with set_forward_context(), pytest.raises(ValueError, match="same global sequence length"):
+            hook.pre_forward(nn.Identity())
+            hook.post_forward(
+                nn.Identity(),
+                (torch.arange(5).reshape(1, 5, 1), torch.arange(6).reshape(1, 6, 1)),
+            )
+
+
 @pytest.mark.cpu
 class TestSequenceParallelConfig:
     """Test SequenceParallelConfig dataclass."""

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -365,8 +365,12 @@ def test_load_weights_maps_mllm_and_preserves_transformer():
     ("parallel_config", "cache_backend", "message"),
     [
         (DiffusionParallelConfig(tensor_parallel_size=2), "none", "Tensor parallelism"),
-        (DiffusionParallelConfig(ulysses_degree=2), "none", "Sequence parallelism"),
-        (DiffusionParallelConfig(ring_degree=2), "none", "Sequence parallelism"),
+        (DiffusionParallelConfig(ring_degree=2), "none", "Ulysses only"),
+        (
+            DiffusionParallelConfig(ulysses_degree=2, cfg_parallel_size=2),
+            "none",
+            "CFG parallelism is not validated",
+        ),
         (
             DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=2),
             "none",
@@ -421,6 +425,96 @@ def test_constructor_accepts_cfg_parallel(mock_dependencies, cfg_parallel_size):
     assert hasattr(pipeline, "predict_noise_with_multi_branch_cfg")
 
 
+def test_constructor_accepts_ulysses_uaa(
+    mock_dependencies,
+    monkeypatch,
+):
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        BooguImagePipeline,
+    )
+
+    monkeypatch.setattr(
+        f"{_MODULE}.current_omni_platform.is_cuda",
+        lambda: True,
+    )
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        tf_model_config=TransformerConfig(
+            params={
+                "num_attention_heads": 28,
+                "num_kv_heads": 7,
+            }
+        ),
+        dtype=torch.float32,
+        parallel_config=DiffusionParallelConfig(
+            ulysses_degree=2,
+            ulysses_mode="advanced_uaa",
+        ),
+    )
+
+    pipeline = BooguImagePipeline(od_config=od_config)
+    assert pipeline.transformer is mock_dependencies["transformer"]
+
+
+def test_constructor_requires_uaa_for_boogu_gqa(
+    mock_dependencies,
+    monkeypatch,
+):
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        BooguImagePipeline,
+    )
+
+    monkeypatch.setattr(
+        f"{_MODULE}.current_omni_platform.is_cuda",
+        lambda: True,
+    )
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        tf_model_config=TransformerConfig(
+            params={
+                "num_attention_heads": 28,
+                "num_kv_heads": 7,
+            }
+        ),
+        dtype=torch.float32,
+        parallel_config=DiffusionParallelConfig(
+            ulysses_degree=2,
+            ulysses_mode="strict",
+        ),
+    )
+
+    with pytest.raises(ValueError, match="advanced_uaa"):
+        BooguImagePipeline(od_config=od_config)
+
+
+def test_constructor_rejects_non_cuda_sequence_parallelism(
+    mock_dependencies,
+    monkeypatch,
+):
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        BooguImagePipeline,
+    )
+
+    monkeypatch.setattr(
+        f"{_MODULE}.current_omni_platform.is_cuda",
+        lambda: False,
+    )
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.float32,
+        parallel_config=DiffusionParallelConfig(
+            ulysses_degree=2,
+            ulysses_mode="advanced_uaa",
+        ),
+    )
+
+    with pytest.raises(NotImplementedError, match="requires CUDA"):
+        BooguImagePipeline(od_config=od_config)
+
+    mock_dependencies["mllm_wrapper"].model.to.assert_not_called()
+
+
 # ---------------------------------------------------------------------------
 # Prompt-encoding tests (deterministic fakes, no constructor)
 # ---------------------------------------------------------------------------

tests/diffusion/models/boogu_image/test_sp_layout.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""L1 unit tests for BOOGU's SP shard-layout arithmetic.
+
+These cover the property the model relies on to build attention masks without
+communicating: every rank can derive *every* rank's segment bounds from the
+global sequence length alone.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.boogu_image.sp_layout import (
+    ShardLayout,
+    rank_concat_mask_or_none,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class TestShardLayout:
+    def test_padding_and_local_extent(self):
+        layout = ShardLayout(original_seq_len=5, world_size=2, rank=0)
+        assert layout.padded_seq_len == 6
+        assert layout.local_seq_len == 3
+        assert layout.padding_size == 1
+        assert layout.bounds(0) == (0, 3)
+        assert layout.bounds(1) == (3, 6)
+
+    def test_valid_lengths_clip_to_each_rank(self):
+        layout = ShardLayout(original_seq_len=5, world_size=2, rank=0)
+        # Sample 0 has 5 valid tokens (3 on rank 0, 2 on rank 1); sample 1 has 2.
+        assert layout.valid_lengths([5, 2], rank=0) == [3, 2]
+        assert layout.valid_lengths([5, 2], rank=1) == [2, 0]
+
+    def test_segment_lengths_split_across_the_boundary(self):
+        layout = ShardLayout(original_seq_len=5, world_size=2, rank=0)
+        # Two reference images of 2 and 3 tokens; the second straddles rank 0/1.
+        assert layout.segment_lengths([[2, 3]], rank=0) == [[2, 1]]
+        assert layout.segment_lengths([[2, 3]], rank=1) == [[0, 2]]
+
+    @pytest.mark.parametrize("world_size", [1, 2, 4])
+    def test_per_rank_lengths_sum_to_the_global_length(self, world_size):
+        layout = ShardLayout(original_seq_len=13, world_size=world_size, rank=0)
+        for global_length in (0, 1, 7, 13):
+            total = sum(layout.valid_lengths([global_length], rank=r)[0] for r in range(world_size))
+            assert total == global_length
+
+
+class TestRankConcatMask:
+    """The mask must describe the sequence Ulysses forms after its all-to-all.
+
+    That sequence is rank0's local shard followed by rank1's, etc. -- not the
+    natural global ordering -- because each rank packs [context|ref|noise]
+    locally before attention.
+    """
+
+    def test_mask_is_concatenation_of_per_rank_segments(self):
+        mask = rank_concat_mask_or_none(
+            [[2], [1]],
+            capacity=3,
+            like=torch.zeros(1, 1),
+        )
+        # rank 0 contributes [T, T, F], rank 1 contributes [T, F, F]
+        assert mask.tolist() == [[True, True, False, True, False, False]]
+
+    def test_all_valid_mask_is_dropped(self):
+        assert (
+            rank_concat_mask_or_none(
+                [[3], [3]],
+                capacity=3,
+                like=torch.zeros(1, 1),
+            )
+            is None
+        )
+
+    def test_required_keeps_an_all_valid_mask(self):
+        mask = rank_concat_mask_or_none(
+            [[3], [3]],
+            capacity=3,
+            like=torch.zeros(1, 1),
+            required=True,
+        )
+        assert mask is not None
+        assert bool(mask.all())
+
+
+class TestRankConcatEquivalence:
+    """The built mask must equal an all-gather of per-rank local masks.
+
+    BOOGU used to hand rank-local masks to the attention layer, which
+    all-gathered them into the global sequence. The model now builds that
+    global mask directly -- saving one collective per attention call -- so the
+    two constructions must agree exactly. This is asserted structurally
+    because the end-to-end SP1-vs-SP2 comparison is numerically insensitive to
+    masking (a run with masks entirely disabled still matches to ~5e-7).
+    """
+
+    @staticmethod
+    def _all_gather_of_local_masks(per_rank_lengths, capacity):
+        """What the deleted `dist.all_gather` path used to produce."""
+        batch_size = len(per_rank_lengths[0])
+        local_masks = []
+        for lengths in per_rank_lengths:
+            local = torch.zeros(batch_size, capacity, dtype=torch.bool)
+            for i, length in enumerate(lengths):
+                local[i, :length] = True
+            local_masks.append(local)
+        return torch.cat(local_masks, dim=1)
+
+    @pytest.mark.parametrize(
+        ("global_len", "world_size", "sample_lengths"),
+        [
+            (5, 2, [5, 2]),
+            (8, 2, [8, 1]),
+            (13, 4, [13, 7, 0]),
+            (16, 4, [16]),
+            (7, 1, [7, 3]),
+        ],
+    )
+    def test_matches_all_gather_of_local_masks(self, global_len, world_size, sample_lengths):
+        layout = ShardLayout(original_seq_len=global_len, world_size=world_size, rank=0)
+        per_rank = [layout.valid_lengths(sample_lengths, rank=r) for r in range(world_size)]
+
+        built = rank_concat_mask_or_none(
+            per_rank,
+            layout.local_seq_len,
+            like=torch.zeros(1, 1),
+            required=True,
+        )
+        assert torch.equal(built, self._all_gather_of_local_masks(per_rank, layout.local_seq_len))