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

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)