[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>

.buildkite/cuda/test-ready.yml

@@ -69,6 +69,7 @@ steps:
         source_file_dependencies: diffusion_distributed_attention
         commands:
           - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda and L4 and cards_2' --run-level "core_model"
+          - timeout 10m pytest -sv tests/diffusion/attention/test_ulysses_uaa.py -m 'core_model and cuda and L4 and cards_2' --run-level "core_model"
         mirror_hardwares: l4_2
 
       - label: "Diffusion · Distributed & Attention Test · 4-GPU"

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))

vllm_omni/diffusion/distributed/sp_plan.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM and The HuggingFace Team
 # Type definitions in this module are adapted from HuggingFace diffusers library:
 #   diffusers/src/diffusers/models/_modeling_parallel.py
@@ -36,7 +37,7 @@ Example:
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeGuard
 
 if TYPE_CHECKING:
     import torch
@@ -220,10 +221,13 @@ class SequenceParallelInput:
             This is useful for layers whose outputs should be split after preprocessing
             (e.g., RoPE embeddings).
         auto_pad: If True, automatically pad the tensor if its size along split_dim
-            is not divisible by world_size. Creates an attention mask to indicate
-            valid vs padding positions. The mask is stored in ForwardContext.
+            is not divisible by world_size. Padding metadata is stored in
+            ForwardContext so models can construct attention masks when needed.
             Note: Ring attention does not support attention mask, so auto_pad
             should only be used with Ulysses SP.
+        shard_group: Optional key shared by tensors representing the same global
+            sequence. Keyed groups track independent padding metadata; omitting
+            it preserves the legacy single-sequence padding behavior.
 
     Example:
         # Split hidden_states along sequence dimension (dim 1)
@@ -240,12 +244,13 @@ class SequenceParallelInput:
     expected_dims: int | None = None
     split_output: bool = False
     auto_pad: bool = False
+    shard_group: str | None = None
 
     def __repr__(self) -> str:
         return (
             f"SequenceParallelInput(split_dim={self.split_dim}, "
             f"expected_dims={self.expected_dims}, split_output={self.split_output}, "
-            f"auto_pad={self.auto_pad})"
+            f"auto_pad={self.auto_pad}, shard_group={self.shard_group!r})"
         )
 
 
@@ -262,6 +267,8 @@ class SequenceParallelOutput:
         gather_dim: The dimension along which to gather the tensor.
         expected_dims: Expected number of dimensions. If provided, validates that
             the tensor has this many dimensions before gathering.
+        shard_group: Optional key whose padding metadata determines the gathered
+            sequence length. Omitting it preserves legacy singleton trimming.
 
     Example:
         # Gather output along sequence dimension (dim 1)
@@ -270,9 +277,13 @@ class SequenceParallelOutput:
 
     gather_dim: int
     expected_dims: int | None = None
+    shard_group: str | None = None
 
     def __repr__(self) -> str:
-        return f"SequenceParallelOutput(gather_dim={self.gather_dim}, expected_dims={self.expected_dims})"
+        return (
+            f"SequenceParallelOutput(gather_dim={self.gather_dim}, "
+            f"expected_dims={self.expected_dims}, shard_group={self.shard_group!r})"
+        )
 
 
 @dataclass(frozen=True)
@@ -372,12 +383,14 @@ SequenceParallelModelPlan = dict[str, SequenceParallelInputType | SequenceParall
 # =============================================================================
 
 
-def _is_valid_input_config(value: object) -> bool:
+def _is_valid_input_config(value: object) -> TypeGuard[AnySequenceParallelInput]:
     """Check if a value is a valid input configuration type."""
     return isinstance(value, (SequenceParallelInput, SequenceParallelPartialInput))
 
 
-def _is_valid_input_config_list(value: object) -> bool:
+def _is_valid_input_config_list(
+    value: object,
+) -> TypeGuard[list[AnySequenceParallelInput] | tuple[AnySequenceParallelInput, ...]]:
     """Check if a value is a list/tuple of valid input configurations."""
     if not isinstance(value, (list, tuple)):
         return False
@@ -396,6 +409,27 @@ def validate_sp_plan(plan: SequenceParallelModelPlan) -> None:
     if not isinstance(plan, dict):
         raise ValueError(f"_sp_plan must be a dict, got {type(plan).__name__}")
 
+    # A shard_group names one split boundary's global sequence. Reusing it
+    # across boundaries would silently overwrite the recorded length.
+    shard_group_owners: dict[str, str] = {}
+
+    def _validate_input(sp_input: AnySequenceParallelInput, *, module_id: str) -> None:
+        if not isinstance(sp_input, SequenceParallelInput) or sp_input.shard_group is None:
+            return
+        if not sp_input.auto_pad:
+            raise ValueError(
+                f"SequenceParallelInput for module {module_id!r} declares "
+                f"shard_group={sp_input.shard_group!r}, but keyed shard "
+                "metadata currently requires auto_pad=True."
+            )
+        owner = shard_group_owners.setdefault(sp_input.shard_group, module_id)
+        if owner != module_id:
+            raise ValueError(
+                f"shard_group={sp_input.shard_group!r} is used by multiple "
+                f"split boundaries ({owner!r} and {module_id!r}). Use a unique "
+                "name for each boundary."
+            )
+
     for module_id, module_plan in plan.items():
         if not isinstance(module_id, str):
             raise ValueError(f"_sp_plan keys must be strings, got {type(module_id).__name__}")
@@ -407,6 +441,8 @@ def validate_sp_plan(plan: SequenceParallelModelPlan) -> None:
             if all(isinstance(x, SequenceParallelOutput) for x in module_plan):
                 continue
             if _is_valid_input_config_list(module_plan):
+                for sp_input in module_plan:
+                    _validate_input(sp_input, module_id=module_id)
                 # List of inputs for a specific parameter (when output is tuple)
                 continue
 
@@ -428,8 +464,10 @@ def validate_sp_plan(plan: SequenceParallelModelPlan) -> None:
                             f"Integer keys (output indices) require split_output=True, "
                             f"got split_output=False for module '{module_id}'[{key}]"
                         )
+                    _validate_input(value, module_id=module_id)
                 elif _is_valid_input_config_list(value):
-                    pass  # Valid list of input configs
+                    for sp_input in value:
+                        _validate_input(sp_input, module_id=module_id)
                 else:
                     raise ValueError(
                         f"Input spec values must be SequenceParallelInput/PartialInput or list thereof, "

vllm_omni/diffusion/forward_context.py

@@ -64,6 +64,10 @@ class ForwardContext:
     sp_padding_size: int = 0
     # Original sequence length before padding (for removing padding in gather)
     sp_original_seq_len: int | None = None
+    # Pre-padding global sequence length per SequenceParallelInput shard_group,
+    # for models that auto-pad several independent sequences in one boundary.
+    # The singleton fields above remain for single-sequence models.
+    sp_shard_metadata: dict[str, int] = field(default_factory=dict)
 
     # Set by registry when _sp_plan hooks are applied.
     # When True, sp_active is determined by _sp_shard_depth (for _sp_plan hooks)
@@ -127,6 +131,13 @@ def is_forward_context_available() -> bool:
     return _forward_context is not None
 
 
+def get_sp_shard_original_seq_len(shard_group: str) -> int | None:
+    """Pre-padding global length of `shard_group`, or None if it was not split."""
+    if not is_forward_context_available():
+        return None
+    return get_forward_context().sp_shard_metadata.get(shard_group)
+
+
 def build_local_sp_padding_mask(
     batch_size: int,
     local_seq_len: int,

vllm_omni/diffusion/hooks/sequence_parallel.py

@@ -182,6 +182,12 @@ class SequenceParallelSplitHook(ModelHook):
         self.module_forward_metadata: ModuleForwardMetadata | None = None
         # Cache for text lengths resolved from kwargs
         self._text_len_cache: dict[str, int] = {}
+        self._shard_groups = {
+            sp_input.shard_group
+            for value in metadata.values()
+            for sp_input in (value if isinstance(value, (list, tuple)) else (value,))
+            if isinstance(sp_input, SequenceParallelInput) and sp_input.shard_group is not None
+        }
 
     def initialize_hook(self, module: nn.Module) -> nn.Module:
         cls = _unwrap_module(module).__class__
@@ -193,6 +199,16 @@ class SequenceParallelSplitHook(ModelHook):
         args_list = list(args)
         # Clear text length cache for this forward pass
         self._text_len_cache.clear()
+        if self._shard_groups:
+            from vllm_omni.diffusion.forward_context import (
+                get_forward_context,
+                is_forward_context_available,
+            )
+
+            if is_forward_context_available():
+                ctx = get_forward_context()
+                for shard_group in self._shard_groups:
+                    ctx.sp_shard_metadata.pop(shard_group, None)
 
         for name, spm in self.metadata.items():
             # Skip if this is a split_output entry (handled in post_forward)
@@ -376,7 +392,11 @@ class SequenceParallelSplitHook(ModelHook):
         if isinstance(sp_input, SequenceParallelInput):
             # Full split with optional auto-padding
             if sp_input.auto_pad:
-                return self._shard_with_auto_pad(x, sp_input.split_dim)
+                return self._shard_with_auto_pad(
+                    x,
+                    sp_input.split_dim,
+                    sp_input.shard_group,
+                )
             _maybe_validate_strict_divisibility(dim=sp_input.split_dim, seq_len=x.size(sp_input.split_dim))
             return sp_shard(x, sp_input.split_dim, validate=False)
         elif isinstance(sp_input, SequenceParallelPartialInput):
@@ -397,13 +417,17 @@ class SequenceParallelSplitHook(ModelHook):
         else:
             raise ValueError(f"Unsupported input config type: {type(sp_input).__name__}")
 
-    def _shard_with_auto_pad(self, x: torch.Tensor, dim: int) -> torch.Tensor:
-        """Shard tensor with automatic padding and attention mask creation.
+    def _shard_with_auto_pad(
+        self,
+        x: torch.Tensor,
+        dim: int,
+        shard_group: str | None,
+    ) -> torch.Tensor:
+        """Pad and shard a tensor while recording its global sequence length.
 
-        When sequence length is not divisible by SP world size, this method:
-        1. Pads the tensor to make it divisible
-        2. Creates an attention mask indicating valid vs padding positions
-        3. Stores the mask and padding info in ForwardContext
+        Models with several independent sequences read the recorded length back
+        via `shard_group` to build masks. Legacy singleton padding fields are
+        also populated for single-sequence models.
         """
         from vllm_omni.diffusion.attention.selector import get_attn_backend_for_capability
         from vllm_omni.diffusion.distributed.parallel_state import (
@@ -418,10 +442,20 @@ class SequenceParallelSplitHook(ModelHook):
             return x
 
         seq_len = x.size(dim)
+        rank = get_sequence_parallel_rank()
         remainder = seq_len % world_size
 
+        if shard_group is not None and is_forward_context_available():
+            ctx = get_forward_context()
+            existing = ctx.sp_shard_metadata.setdefault(shard_group, seq_len)
+            if existing != seq_len:
+                raise ValueError(
+                    f"All tensors in shard_group={shard_group!r} must have the "
+                    f"same global sequence length, but got {existing} and {seq_len}."
+                )
+
         if remainder == 0:
-            # No padding needed
+            # Keyed groups record their length even when no padding is needed.
             return sp_shard(x, dim, validate=False)
 
         # Check backend compatibility
@@ -464,8 +498,11 @@ class SequenceParallelSplitHook(ModelHook):
         x_padded = torch.cat([x, padding], dim=dim)
 
         # Store padding info in forward context (only once, for primary tensor)
-        # Attention layers will create masks dynamically using this info
-        if is_forward_context_available():
+        # Attention layers will create masks dynamically using this info.
+        # Keyed groups record their length in sp_shard_metadata above and must
+        # not clobber the legacy singleton fields, whose semantics assume a
+        # single sequence per boundary.
+        if shard_group is None and is_forward_context_available():
             ctx = get_forward_context()
             # Only set if not already set (first auto_pad tensor wins)
             if ctx.sp_original_seq_len is None:
@@ -477,7 +514,6 @@ class SequenceParallelSplitHook(ModelHook):
                 )
 
         # Shard the padded tensor
-        rank = get_sequence_parallel_rank()
         return x_padded.chunk(world_size, dim=dim)[rank]
 
 
@@ -523,12 +559,6 @@ class SequenceParallelGatherHook(ModelHook):
         if len(output) != len(self.metadata):
             raise ValueError(f"Expected {len(self.metadata)} outputs, got {len(output)}.")
 
-        # Check if padding was applied during split
-        original_seq_len = None
-        if is_forward_context_available():
-            ctx = get_forward_context()
-            original_seq_len = ctx.sp_original_seq_len
-
         actually_gathered = False
 
         for i, spm in enumerate(self.metadata):
@@ -546,6 +576,15 @@ class SequenceParallelGatherHook(ModelHook):
             gathered = sp_gather(x, spm.gather_dim, validate=False)
 
             # Remove padding if it was applied
+            original_seq_len = None
+            if is_forward_context_available():
+                ctx = get_forward_context()
+                if spm.shard_group is None:
+                    original_seq_len = ctx.sp_original_seq_len
+                else:
+                    original_seq_len = ctx.sp_shard_metadata.get(spm.shard_group)
+                    if original_seq_len is None:
+                        raise RuntimeError(f"No SP shard metadata was registered for shard_group={spm.shard_group!r}.")
             if original_seq_len is not None and gathered.size(spm.gather_dim) > original_seq_len:
                 gathered = gathered.narrow(spm.gather_dim, 0, original_seq_len)
                 logger.debug(f"Removed padding: gathered shape {gathered.shape} (original_seq_len={original_seq_len})")
@@ -559,6 +598,9 @@ class SequenceParallelGatherHook(ModelHook):
             ctx._sp_shard_depth = max(0, ctx._sp_shard_depth - 1)
             if ctx._sp_equal_pad_stack:
                 ctx._sp_equal_pad_stack.pop()
+            for spm in self.metadata:
+                if spm is not None and spm.shard_group is not None:
+                    ctx.sp_shard_metadata.pop(spm.shard_group, None)
 
         return output[0] if is_tensor else type(output)(output)
 

vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py

@@ -38,13 +38,22 @@ from vllm.triton_utils import HAS_TRITON
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,
+    SequenceParallelOutput,
+)
 from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
     _fused_cuda_supported,
     fused_qk_norm_rope,
     fused_qk_norm_rope_min_tokens,
 )
 from vllm_omni.diffusion.layers.norm import RMSNorm
-from vllm_omni.diffusion.models.utils import make_attention_mask
+from vllm_omni.diffusion.models.boogu_image.sp_layout import (
+    RotaryEmbedding,
+    ShardLayout,
+    pack_local_rotary,
+    rank_concat_mask_or_none,
+)
 from vllm_omni.platforms import current_omni_platform
 
 if TYPE_CHECKING:
@@ -52,10 +61,6 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
-# (cos, sin) — optionally (cos, sin, packed_table) where packed_table is the
-# fp32 [tokens, head_dim] = [cos(theta) | sin(theta)] layout the fused
-# qk-norm+RoPE op consumes; the eager path ignores the third element.
-RotaryEmbedding = tuple[torch.Tensor, ...]
 RotaryFrequencyTables = list[RotaryEmbedding]
 
 
@@ -540,12 +545,16 @@ def _concat_instruction_image_features(
     assert len(img_tensors) == len(instruct_tensors)
 
     batch_size = img_tensors[0].shape[0]
-    max_seq_len = max(seq_lengths)
+    joint_capacity = instruct_tensors[0].shape[1] + img_tensors[0].shape[1]
 
     concatenated_list = []
     for img_tensor, instruct_tensor in zip(img_tensors, instruct_tensors):
         feature_dim = img_tensor.shape[-1]
-        concatenated = img_tensor.new_zeros(batch_size, max_seq_len, feature_dim)
+        concatenated = img_tensor.new_zeros(
+            batch_size,
+            joint_capacity,
+            feature_dim,
+        )
         for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
             concatenated[i, :encoder_seq_len] = instruct_tensor[i, :encoder_seq_len]
             concatenated[i, encoder_seq_len:seq_len] = img_tensor[i, : seq_len - encoder_seq_len]
@@ -558,16 +567,24 @@ def _split_instruction_image_features(
     hidden_states: torch.Tensor,
     encoder_seq_lengths: list[int],
     seq_lengths: list[int],
+    *,
+    instruct_capacity: int,
+    img_capacity: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Unpack a joint sequence back into (instruction, image) streams."""
     batch_size = hidden_states.shape[0]
     feature_dim = hidden_states.shape[-1]
 
-    max_instruct_len = max(encoder_seq_lengths)
-    max_img_len = max(seq_len - encoder_seq_len for seq_len, encoder_seq_len in zip(seq_lengths, encoder_seq_lengths))
-
-    instruct_hidden_states = hidden_states.new_zeros(batch_size, max_instruct_len, feature_dim)
-    img_hidden_states = hidden_states.new_zeros(batch_size, max_img_len, feature_dim)
+    instruct_hidden_states = hidden_states.new_zeros(
+        batch_size,
+        instruct_capacity,
+        feature_dim,
+    )
+    img_hidden_states = hidden_states.new_zeros(
+        batch_size,
+        img_capacity,
+        feature_dim,
+    )
 
     for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
         img_len = seq_len - encoder_seq_len
@@ -772,7 +789,11 @@ class BooguImageJointAttention(nn.Module):
         attn_output = attn_output.flatten(2, 3).to(dtype)
 
         instruct_attn_out, img_attn_out = _split_instruction_image_features(
-            attn_output, encoder_seq_lengths, seq_lengths
+            attn_output,
+            encoder_seq_lengths,
+            seq_lengths,
+            instruct_capacity=instruct_hidden_states.shape[1],
+            img_capacity=img_hidden_states.shape[1],
         )
         instruct_projected, _ = self.instruct_out(instruct_attn_out)
         img_projected, _ = self.img_out(img_attn_out)
@@ -1139,6 +1160,20 @@ _BOOGU_STACKED_PARAMS_MAPPING = (
 )
 
 
+class _BooguImageSPBoundary(nn.Module):
+    """Identity module used by framework SP split hooks."""
+
+    def forward(self, *tensors: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        return tensors
+
+
+class _BooguImageSPGather(nn.Module):
+    """Identity module used by the framework SP gather hook."""
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return tensor
+
+
 class BooguImageTransformer2DModel(nn.Module):
     """Boogu-Image transformer with mixed stream topology.
 
@@ -1156,6 +1191,46 @@ class BooguImageTransformer2DModel(nn.Module):
         "BooguImageSingleStreamTransformerBlock",
     ]
     _layerwise_offload_blocks_attrs = ["single_stream_layers", "double_stream_layers"]
+    # Boundary outputs are hidden states, masks, and RoPE; each triplet is in
+    # image/reference/context order so tensors sharing a sequence reuse metadata.
+    # Twelve entries: 3 hidden states (dim=3), 3 masks (dim=2), then the
+    # rotary cos/sin pairs for image/reference/context (dim=3 each). Real-valued
+    # RoPE returns tuples, so each rotary embedding contributes two tensors.
+    _sp_plan = {
+        "sp_shard_boundary": {
+            index: SequenceParallelInput(
+                split_dim=1,
+                expected_dims=dims,
+                split_output=True,
+                auto_pad=True,
+                shard_group=shard_group,
+            )
+            for index, (dims, shard_group) in enumerate(
+                zip(
+                    (3, 3, 3, 2, 2, 2, 3, 3, 3, 3, 3, 3),
+                    (
+                        "image",
+                        "reference",
+                        "context",
+                        "image",
+                        "reference",
+                        "context",
+                        "image",
+                        "image",
+                        "reference",
+                        "reference",
+                        "context",
+                        "context",
+                    ),
+                )
+            )
+        },
+        "sp_gather_boundary": SequenceParallelOutput(
+            gather_dim=1,
+            expected_dims=3,
+            shard_group="image",
+        ),
+    }
 
     def __init__(
         self,
@@ -1226,6 +1301,8 @@ class BooguImageTransformer2DModel(nn.Module):
             axes_lens=axes_lens,
             patch_size=patch_size,
         )
+        self.sp_shard_boundary = _BooguImageSPBoundary()
+        self.sp_gather_boundary = _BooguImageSPGather()
 
         self.x_embedder = nn.Linear(
             in_features=patch_size * patch_size * in_channels,
@@ -1472,6 +1549,7 @@ class BooguImageTransformer2DModel(nn.Module):
         l_effective_ref_img_len: list[list[int]],
         l_effective_img_len: list[int],
         temb: torch.Tensor,
+        per_rank_ref_lengths: list[list[list[int]]] | None = None,
     ):
         """Embed image patches and run the refiner blocks.
 
@@ -1481,9 +1559,7 @@ class BooguImageTransformer2DModel(nn.Module):
         avoiding a degenerate zero-length attention.
         """
         batch_size = len(hidden_states)
-        max_combined_img_len = max(
-            img_len + sum(ref_img_len) for img_len, ref_img_len in zip(l_effective_img_len, l_effective_ref_img_len)
-        )
+        max_combined_img_len = hidden_states.shape[1] + ref_image_hidden_states.shape[1]
 
         hidden_states = self.x_embedder(hidden_states)
         ref_image_hidden_states = self.ref_image_patch_embedder(ref_image_hidden_states)
@@ -1502,6 +1578,14 @@ class BooguImageTransformer2DModel(nn.Module):
         flat_l_effective_ref_img_len = list(itertools.chain(*l_effective_ref_img_len))
         num_ref_images = len(flat_l_effective_ref_img_len)
         max_ref_img_len = max(flat_l_effective_ref_img_len)
+        sp_sharded = per_rank_ref_lengths is not None and len(per_rank_ref_lengths) > 1
+        if sp_sharded:
+            # Keep the packed reference batch at the padded local capacity so
+            # every SP rank retains the equal-length contract.
+            max_ref_img_len = max(
+                max_ref_img_len,
+                padded_ref_img_mask.shape[1],
+            )
 
         if max_ref_img_len > 0:
             batch_ref_img_mask = ref_image_hidden_states.new_zeros(num_ref_images, max_ref_img_len, dtype=torch.bool)
@@ -1544,9 +1628,27 @@ class BooguImageTransformer2DModel(nn.Module):
                 # table; the forward-level tuples do not flow into this batch.
                 batch_ref_img_rotary_emb = _with_packed_rope_table(batch_ref_img_rotary_emb)
 
+            if sp_sharded:
+                # Ulysses will concatenate each rank's packed reference batch,
+                # so the mask must cover all ranks' segments.
+                ref_attention_mask = rank_concat_mask_or_none(
+                    [list(itertools.chain(*rank_lengths)) for rank_lengths in per_rank_ref_lengths],
+                    max_ref_img_len,
+                    like=batch_ref_img_mask,
+                )
+            else:
+                ref_attention_mask = (
+                    batch_ref_img_mask
+                    if any(length != max_ref_img_len for length in flat_l_effective_ref_img_len)
+                    else None
+                )
+
             for layer in self.ref_image_refiner:
                 batch_ref_image_hidden_states = layer(
-                    batch_ref_image_hidden_states, batch_ref_img_mask, batch_ref_img_rotary_emb, batch_temb
+                    batch_ref_image_hidden_states,
+                    ref_attention_mask,
+                    batch_ref_img_rotary_emb,
+                    batch_temb,
                 )
 
             # Restore reference-image sequence layout.
@@ -1614,11 +1716,11 @@ class BooguImageTransformer2DModel(nn.Module):
             context_rotary_emb,
             ref_img_rotary_emb,
             noise_rotary_emb,
-            rotary_emb,
-            encoder_seq_lengths,
-            seq_lengths,
-            combined_img_rotary_emb,
-            combined_img_seq_lengths,
+            _,
+            global_encoder_seq_lengths,
+            _,
+            _,
+            _,
         ) = self.rope_embedder(
             freqs_real,
             instruction_attention_mask,
@@ -1629,10 +1731,91 @@ class BooguImageTransformer2DModel(nn.Module):
             device,
         )
 
+        # RoPE is real-valued (cos, sin) tuples; unpack them to shard cos and
+        # sin independently through the split hook, then re-pair below.
+        noise_cos, noise_sin = noise_rotary_emb
+        ref_cos, ref_sin = ref_img_rotary_emb
+        context_cos, context_sin = context_rotary_emb
+        (
+            hidden_states,
+            ref_image_hidden_states,
+            instruction_hidden_states,
+            img_mask,
+            ref_img_mask,
+            instruction_attention_mask,
+            noise_cos,
+            noise_sin,
+            ref_cos,
+            ref_sin,
+            context_cos,
+            context_sin,
+        ) = self.sp_shard_boundary(
+            hidden_states,
+            ref_image_hidden_states,
+            instruction_hidden_states,
+            img_mask,
+            ref_img_mask,
+            instruction_attention_mask,
+            noise_cos,
+            noise_sin,
+            ref_cos,
+            ref_sin,
+            context_cos,
+            context_sin,
+        )
+        noise_rotary_emb = (noise_cos, noise_sin)
+        ref_img_rotary_emb = (ref_cos, ref_sin)
+        context_rotary_emb = (context_cos, context_sin)
+
+        image_layout = ShardLayout.resolve("image", local_seq_len=img_mask.shape[1])
+        reference_layout = ShardLayout.resolve("reference", local_seq_len=ref_img_mask.shape[1])
+        context_layout = ShardLayout.resolve("context", local_seq_len=instruction_attention_mask.shape[1])
+        sp_world_size = image_layout.world_size
+        sp_rank = image_layout.rank
+
+        context_mask_required = context_layout.padding_size > 0 or any(
+            length != context_layout.original_seq_len for length in global_encoder_seq_lengths
+        )
+        noise_mask_required = image_layout.padding_size > 0 or any(
+            length != image_layout.original_seq_len for length in l_effective_img_len
+        )
+        ref_mask_required = reference_layout.padding_size > 0 or any(
+            sum(lengths) != reference_layout.original_seq_len for lengths in l_effective_ref_img_len
+        )
+
+        global_img_lengths = list(l_effective_img_len)
+        # Ulysses concatenates the rank-local sequences after its all-to-all, so
+        # attention masks must describe that rank-concatenated layout. Every
+        # rank's shard bounds are pure arithmetic over the global lengths, so
+        # each rank derives all ranks' segments without any communication.
+        ranks = range(sp_world_size)
+        per_rank_img_lengths = [image_layout.valid_lengths(global_img_lengths, rank=r) for r in ranks]
+        per_rank_ref_lengths = [reference_layout.segment_lengths(l_effective_ref_img_len, rank=r) for r in ranks]
+        per_rank_encoder_lengths = [context_layout.valid_lengths(global_encoder_seq_lengths, rank=r) for r in ranks]
+
+        local_img_lengths = per_rank_img_lengths[sp_rank]
+        local_ref_lengths = per_rank_ref_lengths[sp_rank]
+        local_encoder_lengths = per_rank_encoder_lengths[sp_rank]
+        (
+            rotary_emb,
+            combined_img_rotary_emb,
+            seq_lengths,
+            combined_img_seq_lengths,
+        ) = pack_local_rotary(
+            context_rotary_emb,
+            ref_img_rotary_emb,
+            noise_rotary_emb,
+            local_encoder_lengths,
+            local_ref_lengths,
+            local_img_lengths,
+        )
+
         if _FUSED_QK_NORM_ROPE:
             # One packed [cos|sin] table per rotary embedding that reaches an
             # attention, built once per forward; the tuples grow a third
             # element that the fused path consumes and the eager path ignores.
+            # Under SP, cos/sin are already the rank-local sharded tensors, so
+            # the packed table matches the fused kernel's local Q/K geometry.
             # (ref_img_rotary_emb is packed inside img_patch_embed_and_refine,
             # on the rebuilt per-reference-image batch tuple.)
             context_rotary_emb = _with_packed_rope_table(context_rotary_emb)
@@ -1640,32 +1823,68 @@ class BooguImageTransformer2DModel(nn.Module):
             rotary_emb = _with_packed_rope_table(rotary_emb)
             combined_img_rotary_emb = _with_packed_rope_table(combined_img_rotary_emb)
 
+        context_attention_mask = rank_concat_mask_or_none(
+            per_rank_encoder_lengths,
+            context_layout.local_seq_len,
+            like=instruction_attention_mask,
+            required=context_mask_required,
+        )
+        noise_attention_mask = rank_concat_mask_or_none(
+            per_rank_img_lengths,
+            image_layout.local_seq_len,
+            like=img_mask,
+            required=noise_mask_required,
+        )
+
         # Context refinement.
         for layer in self.context_refiner:
-            instruction_hidden_states = layer(instruction_hidden_states, instruction_attention_mask, context_rotary_emb)
+            instruction_hidden_states = layer(
+                instruction_hidden_states,
+                context_attention_mask,
+                context_rotary_emb,
+            )
 
         # Image patch embedding and refinement.
         combined_img_hidden_states = self.img_patch_embed_and_refine(
             hidden_states,
             ref_image_hidden_states,
-            img_mask,
+            noise_attention_mask,
             ref_img_mask,
             noise_rotary_emb,
             ref_img_rotary_emb,
-            l_effective_ref_img_len,
-            l_effective_img_len,
+            local_ref_lengths,
+            local_img_lengths,
             temb,
+            per_rank_ref_lengths=per_rank_ref_lengths,
         )
 
         instruct_hidden_states = instruction_hidden_states
         img_hidden_states = combined_img_hidden_states
 
-        # Joint mask for [instruct + image].
-        joint_attention_mask = make_attention_mask(hidden_states, seq_lengths)
+        # Joint mask for [instruct + image], over the rank-concatenated sequence.
+        per_rank_combined_img_lengths = [
+            [sum(ref_lengths) + img_length for ref_lengths, img_length in zip(ref_per_sample, img_per_sample)]
+            for ref_per_sample, img_per_sample in zip(per_rank_ref_lengths, per_rank_img_lengths)
+        ]
+        per_rank_seq_lengths = [
+            [encoder_length + combined_length for encoder_length, combined_length in zip(enc_per_sample, comb)]
+            for enc_per_sample, comb in zip(per_rank_encoder_lengths, per_rank_combined_img_lengths)
+        ]
+        joint_attention_mask = rank_concat_mask_or_none(
+            per_rank_seq_lengths,
+            rotary_emb[0].shape[1],
+            like=hidden_states,
+            required=(context_mask_required or noise_mask_required or ref_mask_required),
+        )
 
         # Dual-stream (double-stream) stage.
         if self.num_double_stream_layers > 0:
-            img_attention_mask = make_attention_mask(hidden_states, combined_img_seq_lengths)
+            img_attention_mask = rank_concat_mask_or_none(
+                per_rank_combined_img_lengths,
+                combined_img_rotary_emb[0].shape[1],
+                like=hidden_states,
+                required=noise_mask_required or ref_mask_required,
+            )
 
             for layer in self.double_stream_layers:
                 img_hidden_states, instruct_hidden_states = layer(
@@ -1676,13 +1895,17 @@ class BooguImageTransformer2DModel(nn.Module):
                     combined_img_rotary_emb,
                     rotary_emb,
                     temb,
-                    encoder_seq_lengths,
+                    local_encoder_lengths,
                     seq_lengths,
                 )
 
         # Fuse streams to joint sequence.
-        joint_hidden_states = hidden_states.new_zeros(batch_size, max(seq_lengths), self.hidden_size)
-        for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
+        joint_hidden_states = hidden_states.new_zeros(
+            batch_size,
+            rotary_emb[0].shape[1],
+            self.hidden_size,
+        )
+        for i, (encoder_seq_len, seq_len) in enumerate(zip(local_encoder_lengths, seq_lengths)):
             joint_hidden_states[i, :encoder_seq_len] = instruct_hidden_states[i, :encoder_seq_len]
             joint_hidden_states[i, encoder_seq_len:seq_len] = img_hidden_states[i, : seq_len - encoder_seq_len]
 
@@ -1694,12 +1917,23 @@ class BooguImageTransformer2DModel(nn.Module):
         # Output projection.
         hidden_states = self.norm_out(hidden_states, temb)
 
+        # Extract this rank's noise-image shard and gather it once.
+        local_img_output = hidden_states.new_zeros(
+            batch_size,
+            img_mask.shape[1],
+            hidden_states.shape[-1],
+        )
+        for i, (img_len, seq_len) in enumerate(zip(local_img_lengths, seq_lengths)):
+            local_img_output[i, :img_len] = hidden_states[i, seq_len - img_len : seq_len]
+
+        image_tokens = self.sp_gather_boundary(local_img_output)
+
         # Reshape back to image format.
         p = self.patch_size
         output = []
-        for i, (img_size, img_len, seq_len) in enumerate(zip(img_sizes, l_effective_img_len, seq_lengths)):
+        for i, (img_size, img_len) in enumerate(zip(img_sizes, global_img_lengths)):
             height, width = img_size
-            img_tokens = hidden_states[i][seq_len - img_len : seq_len]
+            img_tokens = image_tokens[i, :img_len]
             img_output = rearrange(
                 img_tokens,
                 "(h w) (p1 p2 c) -> c (h p1) (w p2)",

vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py

@@ -62,6 +62,7 @@ from vllm_omni.diffusion.models.utils import create_transformers_model_with_vllm
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+from vllm_omni.platforms import current_omni_platform
 from vllm_omni.quantization.component_config import resolve_component_quant_config
 
 logger = init_logger(__name__)
@@ -382,7 +383,36 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
         if parallel_config.tensor_parallel_size > 1:
             raise NotImplementedError("Tensor parallelism is not supported by BooguImagePipeline.")
         if (parallel_config.sequence_parallel_size or 1) > 1:
-            raise NotImplementedError("Sequence parallelism is not supported by BooguImagePipeline.")
+            if parallel_config.cfg_parallel_size > 1:
+                # SP × CFG-Parallel is unvalidated: the combined path runs
+                # without crashing but drifts far from the SP=1 baseline
+                # (~19 dB PSNR at 4 steps), so guidance semantics differ.
+                # Remove this guard only alongside a parity test.
+                raise NotImplementedError(
+                    "BooguImagePipeline: sequence parallelism combined with CFG parallelism is not validated."
+                )
+            if parallel_config.ring_degree > 1 or getattr(parallel_config, "allgather_degree", 1) > 1:
+                raise NotImplementedError(
+                    "BooguImagePipeline currently supports sequence parallelism through Ulysses only."
+                )
+            if not current_omni_platform.is_cuda():
+                raise NotImplementedError("BooguImagePipeline sequence parallelism currently requires CUDA.")
+            if parallel_config.ulysses_mode == "strict":
+                # Subscript rather than get(): a missing key means the config
+                # schema has drifted and we would silently validate against a
+                # stale default, not the actual head count.
+                tf_params = self.od_config.tf_model_config.params
+                num_heads = tf_params["num_attention_heads"]
+                num_kv_heads = tf_params["num_kv_heads"]
+                if (
+                    num_heads % parallel_config.ulysses_degree != 0
+                    or num_kv_heads % parallel_config.ulysses_degree != 0
+                ):
+                    raise ValueError(
+                        "BooguImagePipeline GQA heads are not divisible by the "
+                        "configured Ulysses degree; set "
+                        "parallel_config.ulysses_mode='advanced_uaa'."
+                    )
         if parallel_config.use_hsdp:
             raise NotImplementedError("HSDP is not supported by BooguImagePipeline.")
         if self.od_config.cache_backend not in (None, "", "none"):

vllm_omni/diffusion/models/boogu_image/sp_layout.py

@@ -0,0 +1,188 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Shard-layout arithmetic for BOOGU Image sequence parallelism.
+
+BOOGU splits three independent sequences (noise image, reference images, and
+instruction context) at one SP boundary. The framework records each sequence's
+pre-padding global length under a ``shard_group`` key; everything else -- how
+that global length maps onto a given rank's shard -- is arithmetic, so every
+rank can derive the layout of *all* ranks without communicating.
+
+That property is what lets the model build attention masks for the global,
+rank-concatenated sequence that Ulysses produces after its all-to-all, instead
+of all-gathering rank-local masks inside the attention layer.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from vllm_omni.diffusion.forward_context import get_sp_shard_original_seq_len
+
+# (cos, sin) -- optionally (cos, sin, packed_table) where packed_table is the
+# fp32 [tokens, head_dim] = [cos(theta) | sin(theta)] layout the fused
+# qk-norm+RoPE op consumes; the eager path ignores the third element.
+RotaryEmbedding = tuple[torch.Tensor, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class ShardLayout:
+    """How one globally padded sequence maps onto the SP ranks."""
+
+    original_seq_len: int
+    world_size: int
+    rank: int
+
+    @classmethod
+    def resolve(cls, shard_group: str, *, local_seq_len: int) -> ShardLayout:
+        """Read a boundary's layout from the ForwardContext.
+
+        Falls back to an unsharded layout (world_size=1) when SP is off or the
+        tensor was never split, in which case `local_seq_len` is the whole
+        sequence.
+        """
+        from vllm_omni.diffusion.distributed.parallel_state import (
+            get_sequence_parallel_rank,
+            get_sequence_parallel_world_size,
+        )
+
+        original_seq_len = get_sp_shard_original_seq_len(shard_group)
+        if original_seq_len is None:
+            return cls(original_seq_len=local_seq_len, world_size=1, rank=0)
+        return cls(
+            original_seq_len=original_seq_len,
+            world_size=get_sequence_parallel_world_size(),
+            rank=get_sequence_parallel_rank(),
+        )
+
+    @property
+    def padded_seq_len(self) -> int:
+        ws = self.world_size
+        return ((self.original_seq_len + ws - 1) // ws) * ws
+
+    @property
+    def local_seq_len(self) -> int:
+        return self.padded_seq_len // self.world_size
+
+    @property
+    def padding_size(self) -> int:
+        return self.padded_seq_len - self.original_seq_len
+
+    def bounds(self, rank: int) -> tuple[int, int]:
+        """Half-open [start, end) span of the global sequence owned by `rank`."""
+        start = rank * self.local_seq_len
+        return start, start + self.local_seq_len
+
+    def valid_lengths(self, global_lengths: list[int], *, rank: int) -> list[int]:
+        """Per-sample valid prefix lengths clipped to `rank`'s shard."""
+        start, end = self.bounds(rank)
+        return [max(0, min(int(length), end) - start) for length in global_lengths]
+
+    def segment_lengths(self, global_segments: list[list[int]], *, rank: int) -> list[list[int]]:
+        """Per-sample contiguous segments intersected with `rank`'s shard."""
+        start, end = self.bounds(rank)
+        local: list[list[int]] = []
+        for sample in global_segments:
+            sample_local: list[int] = []
+            offset = 0
+            for length in sample:
+                segment_end = offset + int(length)
+                sample_local.append(max(0, min(segment_end, end) - max(offset, start)))
+                offset = segment_end
+            local.append(sample_local)
+        return local
+
+
+def rank_concat_mask_or_none(
+    per_rank_lengths: list[list[int]],
+    capacity: int,
+    *,
+    like: torch.Tensor,
+    required: bool = False,
+) -> torch.Tensor | None:
+    """Mask over the rank-concatenated sequence Ulysses builds post all-to-all.
+
+    ``per_rank_lengths[r][i]`` is sample ``i``'s valid length inside rank ``r``'s
+    shard, and every shard has the same ``capacity`` (guaranteed by auto_pad),
+    so the global sequence is ``world_size * capacity`` long. Returns ``None``
+    for an all-valid mask unless a padding contract requires one.
+    """
+    if not required and all(int(length) == capacity for lengths in per_rank_lengths for length in lengths):
+        return None
+
+    batch_size = len(per_rank_lengths[0])
+    mask = like.new_zeros(batch_size, len(per_rank_lengths) * capacity, dtype=torch.bool)
+    for rank, lengths in enumerate(per_rank_lengths):
+        base = rank * capacity
+        for i, length in enumerate(lengths):
+            mask[i, base : base + int(length)] = True
+    return mask
+
+
+def pack_local_rotary(
+    context_rotary_emb: RotaryEmbedding,
+    ref_img_rotary_emb: RotaryEmbedding,
+    noise_rotary_emb: RotaryEmbedding,
+    encoder_seq_lengths: list[int],
+    ref_img_seq_lengths: list[list[int]],
+    img_seq_lengths: list[int],
+) -> tuple[RotaryEmbedding, RotaryEmbedding, list[int], list[int]]:
+    """Pack local RoPE as ``[context, references, noise]``.
+
+    Capacities stay equal across ranks; returned lengths exclude padding.
+    Rotary embeddings arrive as ``(cos, sin)`` pairs -- the same packing is
+    applied component-wise so both halves stay index-aligned.
+    """
+    context_cos, context_sin = context_rotary_emb
+    ref_cos, ref_sin = ref_img_rotary_emb
+    noise_cos, noise_sin = noise_rotary_emb
+
+    combined_img_seq_lengths = [
+        sum(ref_lengths) + img_length
+        for ref_lengths, img_length in zip(
+            ref_img_seq_lengths,
+            img_seq_lengths,
+        )
+    ]
+    seq_lengths = [
+        context_length + combined_length
+        for context_length, combined_length in zip(
+            encoder_seq_lengths,
+            combined_img_seq_lengths,
+        )
+    ]
+
+    def _pack_component(
+        context_t: torch.Tensor, ref_t: torch.Tensor, noise_t: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size = context_t.shape[0]
+        rotary_dim = context_t.shape[-1]
+        combined_img_capacity = ref_t.shape[1] + noise_t.shape[1]
+        combined_img = context_t.new_zeros(batch_size, combined_img_capacity, rotary_dim)
+        rotary = context_t.new_zeros(
+            batch_size,
+            context_t.shape[1] + combined_img_capacity,
+            rotary_dim,
+        )
+        for i, (context_length, ref_lengths, img_length) in enumerate(
+            zip(encoder_seq_lengths, ref_img_seq_lengths, img_seq_lengths)
+        ):
+            ref_length = sum(ref_lengths)
+            combined_length = ref_length + img_length
+            combined_img[i, :ref_length] = ref_t[i, :ref_length]
+            combined_img[i, ref_length:combined_length] = noise_t[i, :img_length]
+            rotary[i, :context_length] = context_t[i, :context_length]
+            rotary[i, context_length : context_length + combined_length] = combined_img[i, :combined_length]
+        return rotary, combined_img
+
+    rotary_cos, combined_img_cos = _pack_component(context_cos, ref_cos, noise_cos)
+    rotary_sin, combined_img_sin = _pack_component(context_sin, ref_sin, noise_sin)
+
+    return (
+        (rotary_cos, rotary_sin),
+        (combined_img_cos, combined_img_sin),
+        seq_lengths,
+        combined_img_seq_lengths,
+    )