[BugFix][Cosmos3] Pad sound latents so video+sound runs under sequence parallelism (#4678)

Signed-off-by: lishunyang12 <125541396+lishunyang12@users.noreply.github.com>
Signed-off-by: lishunyang <lishunyang12@163.com>

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -1422,7 +1422,7 @@ class TestForwardRouting:
         pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
         sound_latents = torch.zeros(1, 3, 4)
         pipeline._resolve_sound_target_samples = lambda *args: (20, 2.0, 10)
-        pipeline._prepare_sound_latents = lambda *args: (sound_latents, 4)
+        pipeline._prepare_sound_latents = lambda *args, **kwargs: (sound_latents, 4)
         pipeline._decode_sound_latents = lambda *args: torch.ones(1, 2, 20)
         output = pipeline.forward(
             make_request_batch(

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -330,6 +330,27 @@ def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch
         )
 
 
+def test_sound_latent_frames_padded_for_sequence_parallel(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.distributed import parallel_state
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    model = object.__new__(transformer_cosmos3.Cosmos3VFMTransformer)
+    model.latent_patch_size = 2
+    vs = (3, 16, 16)
+
+    monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 1)
+    assert model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=97) == 97
+
+    monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 2)
+    assert model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=97) == 98
+    assert model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=98) == 98
+
+    # ulysses=4, with a transfer control folded into the vision base.
+    monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 4)
+    padded = model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=97, num_vision_items=2)
+    assert (2 * 192 + padded) % 4 == 0
+
+
 def test_compute_rope_freqs_places_text_video_action_and_sound_positions() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer