[perf]: use CUDA IPC in multiproc executor to avoid serialization overhead (#1061)

fastvideo/entrypoints/video_generator.py

@@ -9,6 +9,7 @@ diffusion models.
 import math
 import os
 import re
+import threading
 import time
 from copy import deepcopy
 from typing import Any
@@ -31,6 +32,21 @@ from fastvideo.worker.executor import Executor
 logger = init_logger(__name__)
 
 
+def _infer_latent_batch_size(batch: ForwardBatch) -> int:
+    if isinstance(batch.prompt, list):
+        latent_batch_size = len(batch.prompt)
+    elif batch.prompt is not None:
+        latent_batch_size = 1
+    elif batch.prompt_embeds is not None and len(batch.prompt_embeds) > 0:
+        latent_batch_size = batch.prompt_embeds[0].shape[0]
+    else:
+        raise ValueError(
+            "Cannot infer batch size from batch; no prompt or prompt_embeds found"
+        )
+    latent_batch_size *= batch.num_videos_per_prompt
+    return latent_batch_size
+
+
 class VideoGenerator:
     """
     A unified class for generating videos using diffusion models.
@@ -372,8 +388,31 @@ class VideoGenerator:
 
         # Run inference
         start_time = time.perf_counter()
-        output_batch = self.executor.execute_forward(batch, fastvideo_args)
-        samples = output_batch.output
+
+        # Execute forward pass in a new thread for non-blocking tensor allocation
+        result_container = {}
+
+        def execute_forward_thread():
+            result_container['output_batch'] = self.executor.execute_forward(
+                batch, fastvideo_args)
+
+        thread = threading.Thread(target=execute_forward_thread)
+        thread.start()
+        latent_batch_size = _infer_latent_batch_size(batch)
+        samples = torch.empty((latent_batch_size, 3, sampling_param.num_frames,
+                               sampling_param.height, sampling_param.width),
+                              device='cpu',
+                              pin_memory=fastvideo_args.pin_cpu_memory)
+        thread.join()
+
+        output_batch = result_container['output_batch']
+        if output_batch.output.shape == samples.shape:
+            samples.copy_(output_batch.output)
+        else:
+            logger.warning(
+                "Output shape %s does not match expected shape %s; use slow path",
+                output_batch.output.shape, samples.shape)
+            samples = output_batch.output.cpu()
         logging_info = output_batch.logging_info
 
         gen_time = time.perf_counter() - start_time

fastvideo/pipelines/stages/decoding.py

@@ -242,8 +242,8 @@ class DecodingStage(PipelineStage):
                 decoded_frames = self.decode(cur_latent, fastvideo_args)
                 batch.trajectory_decoded.append(decoded_frames.cpu().float())
 
-        # Convert to CPU float32 for compatibility
-        frames = frames.cpu().float()
+        # Convert to float32 for compatibility
+        frames = frames.to(torch.float32)
 
         # Crop padding if this is a LongCat refinement
         if hasattr(batch, 'num_cond_frames_added') and hasattr(

fastvideo/training/distillation_pipeline.py

@@ -1267,7 +1267,7 @@ class DistillationPipeline(TrainingPipeline):
                     with torch.no_grad():
                         output_batch = self.validation_pipeline.forward(
                             batch, training_args)
-                    samples = output_batch.output
+                    samples = output_batch.output.cpu()
                     if self.rank_in_sp_group != 0:
                         continue
 

fastvideo/training/training_pipeline.py

@@ -857,7 +857,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 # Run validation inference
                 output_batch = self.validation_pipeline.forward(
                     batch, training_args)
-                samples = output_batch.output
+                samples = output_batch.output.cpu()
 
                 if self.rank_in_sp_group != 0:
                     continue

fastvideo/worker/multiproc_executor.py

@@ -650,8 +650,10 @@ class WorkerMultiprocProc:
                         logging_info = None
                         if envs.FASTVIDEO_STAGE_LOGGING:
                             logging_info = output_batch.logging_info
+                        # result tensor shared by CUDA IPC to avoid serialization overhead
+                        result = output_batch.output
                         self.pipe.send({
-                            "output_batch": output_batch.output.cpu(),
+                            "output_batch": result,
                             "logging_info": logging_info,
                             "extra": output_batch.extra,
                         })