[Bugfix][Cosmos3] Auto-pad non-divisible Ulysses GEN sequences (#6918)

Signed-off-by: Rahul Steiger <rsteiger@nvl72d188-T17.cm.cluster>
Co-authored-by: Rahul Steiger <rsteiger@nvl72d188-T17.cm.cluster>

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:

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

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)