[diffusion][feature] Add component-selective offload policies (#5929)

Signed-off-by: lishunyang12 <lishunyang12@163.com>

vllm_omni/diffusion/diffusion_engine.py

@@ -41,6 +41,7 @@ from vllm_omni.diffusion.io_support import (
     supports_audio_output,
     supports_multimodal_input,
 )
+from vllm_omni.diffusion.offloader.config import any_selected_component_uses_allgather
 from vllm_omni.diffusion.output_formatter import (
     format_diffusion_outputs,
     format_empty_diffusion_outputs,
@@ -181,11 +182,7 @@ def _max_num_seqs(od_config: OmniDiffusionConfig) -> int:
 def _uses_dlo_dp_concurrency(od_config: OmniDiffusionConfig) -> bool:
     parallel_config = getattr(od_config, "parallel_config", None)
     dp_size = getattr(parallel_config, "data_parallel_size", 1)
-    return (
-        dp_size > 1
-        and getattr(od_config, "enable_distributed_layerwise_offload", False)
-        and getattr(od_config, "dlo_use_allgather", True)
-    )
+    return dp_size > 1 and any_selected_component_uses_allgather(od_config)
 
 
 def _move_tensor_tree_to_cpu(value: object) -> object:
@@ -1012,33 +1009,12 @@ class DiffusionEngine:
                 raise RuntimeError(f"Could not {action} profiler: {e}") from e
 
     def run_startup_warmup(self) -> None:
-        dlo_use_allgather = getattr(self.od_config, "dlo_use_allgather", True)
-        # Skip dummy run when AllGather is used with more than 1 rank,
-        # because the dummy run sends only 1 request but AllGather requires
-        # all ranks to participate simultaneously.  This covers both DP > 1
-        # and SP > 1 (where dp_size is derived from sp_size in OffloadConfig).
-        pc = getattr(self.od_config, "parallel_config", None)
-        dp_size = getattr(pc, "data_parallel_size", 1) if pc else 1
-        sp_size = getattr(pc, "sequence_parallel_size", 1) if pc else 1
-        effective_shard_size = max(dp_size, sp_size)
-        skip_dummy = (
-            getattr(self.od_config, "enable_distributed_layerwise_offload", False)
-            and dlo_use_allgather
-            and effective_shard_size > 1
-        )
-        if skip_dummy:
-            logger.info(
-                "Skipping dummy run (dist_offload with AllGather, dp_size=%d, sp_size=%d)",
-                dp_size,
-                sp_size,
-            )
-            return
         try:
             self._dummy_run()
         except Exception as e:
             logger.error(f"Dummy run failed: {e}")
             self.close()
-            raise e
+            raise
 
     def _make_dummy_request(
         self,
@@ -1050,10 +1026,13 @@ class DiffusionEngine:
         num_inference_steps: int = 1,
     ) -> OmniDiffusionRequest | None:
         """Build a minimal model request for startup profiling or warmup."""
-        prompt: OmniTextPrompt = {"prompt": "dummy run"}
+        prompt = OmniTextPrompt(prompt="dummy run")
+        model_class_name = self.od_config.model_class_name
+        if model_class_name is None:
+            raise RuntimeError("Dummy request requires a resolved model_class_name")
         supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config)
         if supports_image_input:
-            color_format = image_color_format(self.od_config.model_class_name)
+            color_format = image_color_format(model_class_name)
             images = [PIL.Image.new(color_format, (width, height)) for _ in range(num_image_inputs)]
             prompt.setdefault("multi_modal_data", {})["image"] = images[0] if len(images) == 1 else images
 
@@ -1061,7 +1040,7 @@ class DiffusionEngine:
             audio_sr = 16000
             prompt.setdefault("multi_modal_data", {})["audio"] = np.random.randn(audio_sr * 2).astype(np.float32)
 
-        num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input)
+        num_frames = get_dummy_run_num_frames(model_class_name, supports_audio_input)
         if num_frames <= 0:
             return None
         return OmniDiffusionRequest(
@@ -1101,11 +1080,14 @@ class DiffusionEngine:
             is not DiffusionKVCacheMode.PAGED_SCHEDULER
         ):
             return None
+        model_class_name = self.od_config.model_class_name
+        if model_class_name is None:
+            raise RuntimeError("Diffusion KV profiling requires a resolved model_class_name")
         request = self._make_dummy_request(
             height=1024,
             width=1024,
             guidance_scale=5.0,
-            num_image_inputs=get_dummy_run_num_image_inputs(self.od_config.model_class_name),
+            num_image_inputs=get_dummy_run_num_image_inputs(model_class_name),
         )
         if request is None:
             raise RuntimeError("paged_scheduler requires a runnable Diffusion KV memory profile request")