diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index 2bfa7b558..44bd8e9a7 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -383,6 +383,117 @@ def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None: ) +def test_gen_sp_plan_auto_pads_hidden_states_and_rope_together() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer + + prepare_plan = Cosmos3VFMTransformer._sp_plan["gen_sp_prepare"] + assert set(prepare_plan) == {0, 1, 2} + assert all(spec.auto_pad for spec in prepare_plan.values()) + assert [prepare_plan[index].split_dim for index in range(3)] == [1, 1, 1] + + +def test_gen_sp_prepare_resets_branch_local_padding_metadata() -> None: + from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GenSPPrepare + + ctx = ForwardContext(sp_padding_size=1, sp_original_seq_len=12121) + hidden = torch.zeros(1, 4, 8) + rope = torch.zeros(1, 4, 1, 2) + with override_forward_context(ctx): + output = Cosmos3GenSPPrepare()(hidden, rope, rope) + + assert output[0] is hidden + assert output[1] is rope + assert output[2] is rope + assert ctx.sp_padding_size == 0 + assert ctx.sp_original_seq_len is None + + +def test_sp_attention_masks_padded_gen_and_joint_und_keys() -> None: + from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3CrossAttention + + class RecordingAttention(nn.Module): + def __init__(self) -> None: + super().__init__() + self.metadata: Any | None = None + + def forward(self, query, key, value, metadata): + del key, value + self.metadata = metadata + return query + + module = object.__new__(Cosmos3CrossAttention) + nn.Module.__init__(module) + module.num_heads_local = 2 + module.head_dim = 4 + recording_attention = RecordingAttention() + module.attn = recording_attention + + q = torch.zeros(1, 2, 2, 4) + k = torch.zeros_like(q) + v = torch.zeros_like(q) + k_und = torch.zeros(1, 3, 2, 4) + v_und = torch.zeros_like(k_und) + ctx = ForwardContext(sp_padding_size=1, sp_original_seq_len=3) + with override_forward_context(ctx): + output = module._forward_sp(q, k, v, k_und, v_und) + + assert output.shape == (1, 2, 8) + assert recording_attention.metadata is not None + assert recording_attention.metadata.attn_mask.tolist() == [[True, True, True, False]] + assert recording_attention.metadata.joint_attn_mask.tolist() == [[True, True, True]] + + +def test_ulysses_2d_mask_uses_key_length_with_kv_only_joint_context( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata + from vllm_omni.diffusion.attention.parallel import ulysses + + def fake_all_to_all( + _group: object, + tensor: torch.Tensor, + _scatter_idx: int, + _gather_idx: int, + _use_sync: bool, + ) -> torch.Tensor: + return torch.cat([tensor, tensor], dim=1)[:, :, :1] + + monkeypatch.setattr(ulysses.SeqAllToAll4D, "apply", staticmethod(fake_all_to_all)) + strategy = ulysses.UlyssesParallelAttention( + sp_group=SimpleNamespace( + ulysses_group=object(), + ulysses_world_size=2, + ulysses_rank=0, + ring_world_size=1, + ), + scatter_idx=2, + gather_idx=1, + use_sync=False, + ) + query = torch.zeros(1, 2, 2, 4) + key = torch.zeros_like(query) + value = torch.zeros_like(query) + joint_query = torch.zeros(1, 0, 2, 4) + joint_key = torch.zeros(1, 3, 2, 4) + metadata = AttentionMetadata( + attn_mask=torch.tensor([[True, True, True, False]]), + joint_attn_mask=torch.ones(1, 3, dtype=torch.bool), + joint_query=joint_query, + joint_key=joint_key, + joint_value=torch.zeros_like(joint_key), + joint_strategy="front", + ) + + query_out, key_out, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata) + + assert query_out.shape[1] == 4 + assert key_out.shape[1] == 7 + assert metadata_out is not None + assert metadata_out.attn_mask.tolist() == [[True, True, True, True, True, True, False]] + + def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> None: from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 @@ -721,7 +832,7 @@ def test_forward_returns_video_plus_optional_modality_predictions( assert [tuple(tensor.shape) for tensor in output] == expected_shapes -def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch: pytest.MonkeyPatch) -> None: +def test_forward_with_sound_odd_sequence_uses_auto_padding_contract(monkeypatch: pytest.MonkeyPatch) -> None: import vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 as cosmos3_module model = cosmos3_module.Cosmos3VFMTransformer( @@ -732,16 +843,17 @@ def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch ) monkeypatch.setattr(cosmos3_module, "_get_ulysses_state", lambda: (2, 0, None)) - with pytest.raises(ValueError, match=r"GEN sequence length \(3 = video tokens 2 \+ sound tokens 1\)"): - model( - hidden_states=torch.zeros(1, 2, 1, 1, 2), - timestep=torch.tensor([1.0]), - text_ids=torch.tensor([[1, 2]], dtype=torch.long), - text_mask=torch.ones(1, 2, dtype=torch.long), - video_shape=(1, 1, 2), - fps=24.0, - sound_latents=torch.zeros(1, 3, 1), - ) + output = model( + hidden_states=torch.zeros(1, 2, 1, 1, 2), + timestep=torch.tensor([1.0]), + text_ids=torch.tensor([[1, 2]], dtype=torch.long), + text_mask=torch.ones(1, 2, dtype=torch.long), + video_shape=(1, 1, 2), + fps=24.0, + sound_latents=torch.zeros(1, 3, 1), + ) + + assert [tuple(tensor.shape) for tensor in output] == [(1, 2, 1, 1, 2), (1, 3, 1)] def test_sound_latent_frames_padded_for_sequence_parallel(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/vllm_omni/diffusion/attention/parallel/ulysses.py b/vllm_omni/diffusion/attention/parallel/ulysses.py index cbaa8ba62..28298c0a2 100644 --- a/vllm_omni/diffusion/attention/parallel/ulysses.py +++ b/vllm_omni/diffusion/attention/parallel/ulysses.py @@ -608,15 +608,15 @@ class UlyssesParallelAttention: if attn_metadata.attn_mask is None: assert attn_metadata.joint_attn_mask is not None attn_metadata.attn_mask = torch.ones( - [query.shape[0], query.shape[1] - attn_metadata.joint_attn_mask.shape[1]], + [key.shape[0], key.shape[1] - attn_metadata.joint_attn_mask.shape[1]], dtype=torch.bool, - device=query.device, + device=key.device, ) elif attn_metadata.joint_attn_mask is None: attn_metadata.joint_attn_mask = torch.ones( - [query.shape[0], query.shape[1] - attn_metadata.attn_mask.shape[1]], + [key.shape[0], key.shape[1] - attn_metadata.attn_mask.shape[1]], dtype=torch.bool, - device=query.device, + device=key.device, ) attn_metadata.attn_mask = ( torch.cat([attn_metadata.joint_attn_mask, attn_metadata.attn_mask], dim=1) @@ -625,9 +625,10 @@ class UlyssesParallelAttention: ) if attn_metadata.attn_mask is not None: - # the final attn_mask is ready, the length should be aligedn with query length - assert attn_metadata.attn_mask.shape[1] == query.shape[1], ( - f"attn_mask length: {attn_metadata.attn_mask.shape[1]} != query length: {query.shape[1]}" + # A 2D mask describes keys. Joint attention may provide K/V-only + # context, so key and query lengths are not necessarily equal. + assert attn_metadata.attn_mask.shape[1] == key.shape[1], ( + f"attn_mask length: {attn_metadata.attn_mask.shape[1]} != key length: {key.shape[1]}" ) attn_metadata.attn_mask = attn_metadata.attn_mask.bool().contiguous() return query, key, value, attn_metadata, ctx diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index c946eb4ca..fe7a34c8b 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -705,7 +705,21 @@ class Cosmos3CrossAttention(nn.Module): # standard reverse-all-to-all path (no joint-output splitting). joint_q = q.new_empty(B, 0, self.num_heads_local, self.head_dim) + gen_mask = None + joint_mask = None + if is_forward_context_available(): + ctx = get_forward_context() + if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: + padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size + gen_mask = torch.ones(B, padded_seq_len, dtype=torch.bool, device=q.device) + gen_mask[:, ctx.sp_original_seq_len :] = False + # UND K/V are trimmed to the common real text length before GEN + # attention, so every remaining joint key is valid. + joint_mask = torch.ones(B, k_und.shape[1], dtype=torch.bool, device=q.device) + attn_metadata = AttentionMetadata( + attn_mask=gen_mask, + joint_attn_mask=joint_mask, joint_query=joint_q, joint_key=k_und, joint_value=v_und, @@ -1055,6 +1069,13 @@ class Cosmos3GenSPPrepare(nn.Module): freqs_cos: torch.Tensor, freqs_sin: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # A transfer denoising step may run a full-control branch followed by a + # shorter no-control branch in one ForwardContext. Auto-padding metadata + # belongs to one branch only and must not truncate a later gather. + if is_forward_context_available(): + ctx = get_forward_context() + ctx.sp_padding_size = 0 + ctx.sp_original_seq_len = None return hidden_gen, freqs_cos, freqs_sin @@ -1112,9 +1133,9 @@ class Cosmos3VFMTransformer(nn.Module): _sp_plan = { "gen_sp_prepare": { - 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True), - 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), - 2: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True), + 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), + 2: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), }, "gen_sp_gather": SequenceParallelOutput(gather_dim=1, expected_dims=3), } @@ -1586,47 +1607,6 @@ class Cosmos3VFMTransformer(nn.Module): self.cached_kv = None self.cached_freqs_gen = None - @staticmethod - def _validate_gen_sequence_parallel( - *, - s_gen: int, - s_video: int, - s_control: int, - s_action: int, - s_sound: int, - has_action: bool, - has_sound: bool, - has_control: bool, - ulysses_size: int, - ) -> None: - if ulysses_size <= 1 or s_gen % ulysses_size == 0: - return - - detail_parts = [] - if has_control: - detail_parts.append(f"control tokens {s_control}") - detail_parts.append(f"video tokens {s_video}") - if has_action: - detail_parts.append(f"action tokens {s_action}") - if has_sound: - detail_parts.append(f"sound tokens {s_sound}") - detail = " = " + " + ".join(detail_parts) if len(detail_parts) > 1 else "" - adjust_detail = ( - "Adjust the spatial resolution, frame count, action chunk size, " - "sound duration, or sound latent FPS so the combined media sequence is a " - "multiple of ulysses_degree." - if has_control or has_action or has_sound - else ( - "Adjust the spatial resolution so that " - "t * ceil(h/patch) * ceil(w/patch) is a multiple " - "of ulysses_degree." - ) - ) - raise ValueError( - f"GEN sequence length ({s_gen}{detail}) must be divisible by " - f"ulysses_degree ({ulysses_size}). {adjust_detail}" - ) - def sound_latent_frames_for_sequence_parallel( self, *, @@ -1907,17 +1887,6 @@ class Cosmos3VFMTransformer(nn.Module): # framework Attention handles the Ulysses head-slicing internally. if self.cached_kv is None or self.cached_freqs_gen is None: raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.") - self._validate_gen_sequence_parallel( - s_gen=hidden_gen.shape[1], - s_video=s_video, - s_control=s_control, - s_action=s_action, - s_sound=s_sound, - has_action=has_action, - has_sound=has_sound, - has_control=has_control, - ulysses_size=1 if use_multi_control_attention else ulysses_size, - ) freqs_cos, freqs_sin = self.cached_freqs_gen if not use_multi_control_attention: hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin)