[Bugfix][Cosmos3] Support multi-chunk transfer with distributed VAE (#6920)

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

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -1150,6 +1150,72 @@ def test_transfer_fps_matches_resolved_frame_rate_precedence() -> None:
     assert cfg.fps == sp.resolved_frame_rate == 12.0
 
 
+def test_transfer_vae_executor_requires_distributed_vae() -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+
+    pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
+    executor = object()
+    pipeline.vae = SimpleNamespace(distributed_executor=executor, is_distributed_enabled=lambda: True)
+    assert pipeline._transfer_vae_executor() is executor
+
+    pipeline.vae = SimpleNamespace(distributed_executor=executor, is_distributed_enabled=lambda: False)
+    assert pipeline._transfer_vae_executor() is None
+
+
+def test_sync_transfer_overlap_slices_output_rank_and_broadcasts() -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+
+    class RecordingExecutor:
+        rank = 0
+
+        def __init__(self) -> None:
+            self.value = None
+
+        def broadcast_tensor(self, value):
+            self.value = value
+            return value
+
+    pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
+    output = torch.arange(1 * 3 * 5 * 2 * 2).reshape(1, 3, 5, 2, 2)
+    executor = RecordingExecutor()
+
+    overlap = pipeline._sync_transfer_overlap(
+        output,
+        overlap_frames=2,
+        reference_video=output,
+        vae_executor=executor,
+    )
+
+    assert overlap is not None
+    torch.testing.assert_close(overlap, output[:, :, -2:])
+    assert executor.value is overlap
+
+
+def test_sync_transfer_overlap_allocates_receive_buffer_on_non_output_rank() -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+
+    class ReceivingExecutor:
+        rank = 1
+
+        def broadcast_tensor(self, value):
+            assert value.shape == (1, 3, 2, 4, 5)
+            return torch.ones_like(value)
+
+    pipeline = SimpleNamespace(device=torch.device("cpu"), vae=SimpleNamespace(dtype=torch.float32))
+    reference = torch.zeros(1, 3, 6, 4, 5)
+
+    overlap = Cosmos3OmniDiffusersPipeline._sync_transfer_overlap(
+        pipeline,
+        torch.empty(0),
+        overlap_frames=2,
+        reference_video=reference,
+        vae_executor=ReceivingExecutor(),
+    )
+
+    assert overlap is not None
+    assert torch.equal(overlap, torch.ones_like(overlap))
+
+
 def test_transfer_edge_uses_rgb_canny(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transfer
 

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -3215,6 +3215,46 @@ class Cosmos3OmniDiffusersPipeline(
             self.transformer.reset_cache()
         return latents
 
+    def _transfer_vae_executor(self) -> Any | None:
+        """Return the active distributed VAE executor, if any."""
+        executor = getattr(self.vae, "distributed_executor", None)
+        is_distributed_enabled = getattr(self.vae, "is_distributed_enabled", None)
+        if executor is None or not callable(is_distributed_enabled) or not is_distributed_enabled():
+            return None
+        return executor
+
+    def _sync_transfer_overlap(
+        self,
+        output_video: torch.Tensor,
+        *,
+        overlap_frames: int,
+        reference_video: torch.Tensor,
+        vae_executor: Any | None,
+    ) -> torch.Tensor | None:
+        """Broadcast only the decoded frames needed to condition the next chunk."""
+        if overlap_frames <= 0:
+            return None
+
+        is_output_rank = vae_executor is None or vae_executor.rank == 0
+        if is_output_rank:
+            overlap = output_video[:, :, -overlap_frames:].contiguous()
+        else:
+            overlap = torch.empty(
+                (
+                    reference_video.shape[0],
+                    reference_video.shape[1],
+                    overlap_frames,
+                    reference_video.shape[3],
+                    reference_video.shape[4],
+                ),
+                device=self.device,
+                dtype=self.vae.dtype,
+            )
+
+        if vae_executor is not None:
+            overlap = vae_executor.broadcast_tensor(overlap)
+        return overlap
+
     def _forward_transfer(
         self,
         *,
@@ -3354,6 +3394,8 @@ class Cosmos3OmniDiffusersPipeline(
         output_chunks: list[torch.Tensor] = []
         control_chunks_per_hint: dict[str, list[torch.Tensor]] = {key: [] for key in per_hint_frames}
         previous_output: torch.Tensor | None = None
+        vae_executor = self._transfer_vae_executor()
+        is_output_rank = vae_executor is None or vae_executor.rank == 0
 
         for chunk_id in range(num_chunks):
             start_frame = chunk_id * stride
@@ -3448,16 +3490,26 @@ class Cosmos3OmniDiffusersPipeline(
                 generator=generator,
             )
             output_video = self._decode_latents(latents).clamp(-1, 1)
-            previous_output = output_video
+            if chunk_id + 1 < num_chunks:
+                previous_output = self._sync_transfer_overlap(
+                    output_video,
+                    overlap_frames=min(transfer_config.num_conditional_frames, chunk_frames),
+                    reference_video=target_norm,
+                    vae_executor=vae_executor,
+                )
 
-            if chunk_id == 0:
-                output_chunks.append(output_video)
-                for key, control in control_norms.items():
-                    control_chunks_per_hint[key].append(control)
-            else:
-                output_chunks.append(output_video[:, :, current_conditional_frames:])
-                for key, control in control_norms.items():
-                    control_chunks_per_hint[key].append(control[:, :, current_conditional_frames:])
+            if is_output_rank:
+                if chunk_id == 0:
+                    output_chunks.append(output_video)
+                    for key, control in control_norms.items():
+                        control_chunks_per_hint[key].append(control)
+                else:
+                    output_chunks.append(output_video[:, :, current_conditional_frames:])
+                    for key, control in control_norms.items():
+                        control_chunks_per_hint[key].append(control[:, :, current_conditional_frames:])
+
+        if not is_output_rank:
+            return DiffusionOutput(output={"video": output_video}, custom_output={"fps": frame_rate})
 
         full_output = torch.cat(output_chunks, dim=2)[:, :, :total_frames]
         full_controls = {