diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index 633abb9d3..54fed272c 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.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" diff --git a/tests/diffusion/attention/test_ulysses_uaa.py b/tests/diffusion/attention/test_ulysses_uaa.py index 230eb8b71..3782d8dd6 100644 --- a/tests/diffusion/attention/test_ulysses_uaa.py +++ b/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", [ diff --git a/tests/diffusion/distributed/test_boogu_image_sp.py b/tests/diffusion/distributed/test_boogu_image_sp.py new file mode 100644 index 000000000..276744a99 --- /dev/null +++ b/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 diff --git a/tests/diffusion/distributed/test_sp_plan_hooks.py b/tests/diffusion/distributed/test_sp_plan_hooks.py index 9a3c0ee82..e40ec8676 100644 --- a/tests/diffusion/distributed/test_sp_plan_hooks.py +++ b/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.""" diff --git a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py index 4abd7ef9b..9ad022a10 100644 --- a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py +++ b/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) # --------------------------------------------------------------------------- diff --git a/tests/diffusion/models/boogu_image/test_sp_layout.py b/tests/diffusion/models/boogu_image/test_sp_layout.py new file mode 100644 index 000000000..a02cdea68 --- /dev/null +++ b/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)) diff --git a/vllm_omni/diffusion/distributed/sp_plan.py b/vllm_omni/diffusion/distributed/sp_plan.py index b8d248865..3e5568706 100644 --- a/vllm_omni/diffusion/distributed/sp_plan.py +++ b/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, " diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index 55aaf920b..8068fe3ed 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/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, diff --git a/vllm_omni/diffusion/hooks/sequence_parallel.py b/vllm_omni/diffusion/hooks/sequence_parallel.py index 95f5b0f47..32f15de5f 100644 --- a/vllm_omni/diffusion/hooks/sequence_parallel.py +++ b/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) diff --git a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py index f3fdeedc8..4088dd5c9 100644 --- a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py +++ b/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)", diff --git a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py index b4cd930b2..d96d4a6c1 100644 --- a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py +++ b/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"): diff --git a/vllm_omni/diffusion/models/boogu_image/sp_layout.py b/vllm_omni/diffusion/models/boogu_image/sp_layout.py new file mode 100644 index 000000000..feaa457b9 --- /dev/null +++ b/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, + )