[diffusion] fix: fix SANA-WM CFG-parallel tensor devices (#28835)

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/base.py

@@ -838,9 +838,11 @@ class SanaWMDenoisingStage(DenoisingStage):
         scheduler = getattr(
             batch, "scheduler", None
         ) or get_or_create_request_scheduler(batch, self.scheduler)
+        self._move_scheduler_tensors_to_device(scheduler, device)
         timesteps = batch.timesteps
         if timesteps is None:
             raise ValueError("SANA-WM denoising requires prepared timesteps.")
+        timesteps = timesteps.to(device=device)
 
         latents = batch.latents.to(device=device, dtype=target_dtype)
         init_latents = latents.clone()
@@ -1040,6 +1042,17 @@ class SanaWMDenoisingStage(DenoisingStage):
         batch.latents = server_args.pipeline_config.post_denoising_loop(latents, batch)
         return batch
 
+    @staticmethod
+    def _move_scheduler_tensors_to_device(scheduler: object, device) -> None:
+        for name in ("sigmas", "timesteps"):
+            for attr_name in (name, f"_{name}"):
+                value = getattr(scheduler, attr_name, None)
+                if isinstance(value, torch.Tensor):
+                    try:
+                        setattr(scheduler, attr_name, value.to(device=device))
+                    except AttributeError:
+                        pass
+
 
 class SanaWMBeforeDenoisingStage(PipelineStage):
     """

python/sglang/multimodal_gen/test/unit/sana_wm/test_pipeline_config.py

@@ -1079,6 +1079,44 @@ class TestSanaWMDenoisingStage(unittest.TestCase):
         self.assertTrue(torch.allclose(combined_from_pos_rank, serial))
         self.assertTrue(torch.allclose(combined_from_neg_rank, serial))
 
+    def test_scheduler_tensors_move_to_target_device(self) -> None:
+        scheduler = SimpleNamespace(
+            sigmas=torch.tensor([1.0, 0.0]),
+            timesteps=torch.tensor([1000.0, 0.0]),
+            untouched="value",
+        )
+
+        SanaWMDenoisingStage._move_scheduler_tensors_to_device(
+            scheduler, torch.device("cpu")
+        )
+
+        self.assertEqual(scheduler.sigmas.device.type, "cpu")
+        self.assertEqual(scheduler.timesteps.device.type, "cpu")
+        self.assertEqual(scheduler.untouched, "value")
+
+        class PropertyScheduler:
+            def __init__(self):
+                self._sigmas = torch.tensor([1.0, 0.0])
+                self._timesteps = torch.tensor([1000.0, 0.0])
+                self.untouched = "value"
+
+            @property
+            def sigmas(self):
+                return self._sigmas
+
+            @property
+            def timesteps(self):
+                return self._timesteps
+
+        property_scheduler = PropertyScheduler()
+        SanaWMDenoisingStage._move_scheduler_tensors_to_device(
+            property_scheduler, torch.device("cpu")
+        )
+
+        self.assertEqual(property_scheduler.sigmas.device.type, "cpu")
+        self.assertEqual(property_scheduler.timesteps.device.type, "cpu")
+        self.assertEqual(property_scheduler.untouched, "value")
+
 
 class TestSanaWMNativeDiTChunking(unittest.TestCase):
     def test_softmax_chunking_is_disabled_by_default_for_upstream_parity(self) -> None: