[Perf][Cosmos3] Optimize I2V conditioning transfer (8.63% faster x-inference)

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -783,6 +783,38 @@ def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline)
     torch.testing.assert_close(action, clean)
 
 
+def test_i2v_conditioning_moves_single_frame_before_video_expand(make_cosmos3_pipeline) -> None:
+    pipeline = make_cosmos3_pipeline()
+    calls: list[tuple[str, tuple[int, ...]]] = []
+
+    def record_to_vae_device(tensor: torch.Tensor, *, pin_cpu: bool = False) -> torch.Tensor:
+        assert pin_cpu is True
+        calls.append(("to_vae_device", tuple(tensor.shape)))
+        return tensor
+
+    class RecordingVAE(StubCosmos3VAE):
+        def encode(self, video: torch.Tensor):
+            calls.append(("encode", tuple(video.shape)))
+            assert video.is_contiguous()
+            return super().encode(video)
+
+    pipeline._to_vae_device = record_to_vae_device
+    pipeline.vae = RecordingVAE(z_dim=2)
+
+    latent = pipeline._encode_conditioning_video(
+        torch.zeros(1, 3, 16, 24),
+        num_frames=5,
+        height=16,
+        width=24,
+    )
+
+    assert calls == [
+        ("to_vae_device", (1, 3, 16, 24)),
+        ("encode", (1, 3, 5, 16, 24)),
+    ]
+    assert latent.shape == (1, 2, 2, 2, 3)
+
+
 def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> None:
     pipeline = make_cosmos3_pipeline()
     latents = torch.zeros(1, 2, 1, 1, 1)

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -1644,18 +1644,36 @@ def _resolve_sound_target_samples(
 
     # -- VAE decode ----------------------------------------------------------
 
+    def _get_latents_mean_std(self, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
+        cached = getattr(self, "_latents_mean_std", None)
+        if cached is not None:
+            latents_mean, latents_std = cached
+            if latents_mean.device == device and latents_mean.dtype == dtype:
+                return latents_mean, latents_std
+
+        latents_mean = torch.as_tensor(self.vae.config.latents_mean, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
+        latents_std = torch.as_tensor(self.vae.config.latents_std, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
+        self._latents_mean_std = (latents_mean, latents_std)
+        return latents_mean, latents_std
+
+    def _to_vae_device(self, tensor: torch.Tensor, *, pin_cpu: bool = False) -> torch.Tensor:
+        if tensor.device == self.device and tensor.dtype == self.vae.dtype:
+            return tensor
+
+        non_blocking = False
+        if tensor.device.type == "cpu" and self.device.type == "cuda":
+            if pin_cpu and not tensor.is_pinned():
+                tensor = tensor.pin_memory()
+            non_blocking = tensor.is_pinned()
+
+        return tensor.to(device=self.device, dtype=self.vae.dtype, non_blocking=non_blocking)
+
     def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
         latents = latents.to(self.vae.dtype)
 
         if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"):
-            if not hasattr(self, "_latents_mean"):
-                self._latents_mean = (
-                    torch.tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1).to(self.device, self.vae.dtype)
-                )
-                self._latents_std = (
-                    torch.tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1).to(self.device, self.vae.dtype)
-                )
-            latents = (latents * self._latents_std) + self._latents_mean
+            latents_mean, latents_std = self._get_latents_mean_std(latents.device, latents.dtype)
+            latents = (latents * latents_std) + latents_mean
         else:
             scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
             latents = latents / scaling_factor
@@ -1769,18 +1787,17 @@ def _encode_conditioning_video(
         The caller keeps only the conditioned latent frame(s) and replaces
         the rest with noise.
         """
-        # image_tensor: [1, 3, H, W] -> [1, 3, num_frames, H, W]
+        # Move the single conditioning frame first, then materialize the
+        # repeated video on GPU.  Materializing [1, 3, T, H, W] on CPU creates
+        # a large pageable H2D copy on I2V requests.
+        image_tensor = self._to_vae_device(image_tensor, pin_cpu=True)
         video = image_tensor.unsqueeze(2).expand(-1, -1, num_frames, -1, -1).contiguous()
-        video = video.to(device=self.device, dtype=self.vae.dtype)
 
         latent = self.vae.encode(video).latent_dist.mode()
 
         # Normalize (inverse of _decode_latents denormalization)
         if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"):
-            latents_mean = (
-                torch.tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype)
-            )
-            latents_std = torch.tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype)
+            latents_mean, latents_std = self._get_latents_mean_std(latent.device, latent.dtype)
             latent = (latent - latents_mean) / latents_std
         else:
             scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
@@ -1818,14 +1835,11 @@ def _encode_video_tensor(self, video_tensor: torch.Tensor, image_size: Any | Non
         if video_tensor.shape[0] != 1 or video_tensor.shape[1] != 3:
             raise ValueError(f"Cosmos3 video tensor must have shape [1, 3, T, H, W], got {tuple(video_tensor.shape)}.")
 
-        video = video_tensor.to(device=self.device, dtype=self.vae.dtype)
+        video = self._to_vae_device(video_tensor)
         latent = self.vae.encode(video).latent_dist.mode()
 
         if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"):
-            latents_mean = (
-                torch.tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype)
-            )
-            latents_std = torch.tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype)
+            latents_mean, latents_std = self._get_latents_mean_std(latent.device, latent.dtype)
             latent = (latent - latents_mean) / latents_std
         else:
             scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)