[Bugfix] Maintain model-level CPU offload in a blocking way (#1978)

Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>

vllm_omni/diffusion/offloader/sequential_backend.py

@@ -43,9 +43,17 @@ class SequentialOffloadHook(ModelHook):
         module: nn.Module,
         target_device: torch.device,
         *,
-        non_blocking: bool = True,
+        non_blocking: bool = False,
         pin_memory: bool = False,
     ) -> None:
+        """Move module parameters and buffers to device.
+
+        This cls method specifically prevents recursion device movement,
+        E.g., Cache-DiT CachedBlocks has attr `transformer` as a ref to original
+        transformer blocks, thus `module.to(device)` will fail for recursion calling,
+        refer to
+        https://github.com/vipshop/cache-dit/blob/v1.2.3/src/cache_dit/caching/cache_blocks/__init__.py#L83
+        """
         for p in module.parameters():
             if p.data.device != target_device:
                 data = p.data.to(target_device, non_blocking=non_blocking)
@@ -83,7 +91,7 @@ class SequentialOffloadHook(ModelHook):
         except StopIteration:
             return
 
-        self._move_params(module, self.device)
+        self._move_params(module, self.device, non_blocking=False)
 
     def pre_forward(self, module: nn.Module, *args, **kwargs) -> tuple[tuple, dict]:
         # Offload target modules to CPU