[Diffusion] Fix Cosmos3 serving and rollout resource cleanup

python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py

@@ -17,7 +17,7 @@
 import pprint
 from collections import Counter
 from copy import deepcopy
-from dataclasses import MISSING, asdict, dataclass, field, fields
+from dataclasses import MISSING, dataclass, field, fields
 from typing import TYPE_CHECKING, Any, Optional, Sequence, Union
 
 import PIL.Image
@@ -413,7 +413,10 @@ def validate(self):
         self.metrics = RequestMetrics(request_id=self.request_id)
 
     def __str__(self):
-        return pprint.pformat(asdict(self), indent=2, width=120)
+        # Logging must not deepcopy live tensors or rollout controller state.
+        values = {f.name: getattr(self, f.name) for f in fields(self)}
+        values["rollout_controllers"] = list(self.rollout_controllers)
+        return pprint.pformat(values, indent=2, width=120)
 
     def log(self, server_args: ServerArgs):
         if self.is_warmup or self.suppress_logs:

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -160,6 +160,7 @@
     finalize_rollout_if_enabled,
     init_rollout_env_if_enabled,
     prepare_rollout_if_enabled,
+    release_rollout_resources,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -2042,8 +2043,11 @@ def forward(
         batch: Req,
         server_args: ServerArgs,
     ) -> Req:
-        with self._offload_for_torch_compile_warmup(batch):
-            return self._denoise(batch, server_args)
+        try:
+            with self._offload_for_torch_compile_warmup(batch):
+                return self._denoise(batch, server_args)
+        finally:
+            release_rollout_resources(batch)
 
     @torch.no_grad()
     def _denoise(
@@ -2138,7 +2142,7 @@ def _denoise(
             batch=batch,
             latents=ctx.latents,
             num_inference_steps=num_timesteps,
-            final_timestep=timesteps_cpu.new_zeros(()),
+            final_timestep=ctx.timesteps.new_zeros(()),
             pipeline_config=server_args.pipeline_config,
         )
         self._finalize_denoising_loop(ctx, batch, server_args)

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py

@@ -66,6 +66,7 @@
     finalize_rollout_if_enabled,
     init_rollout_env_if_enabled,
     prepare_rollout_if_enabled,
+    release_rollout_resources,
 )
 from sglang.multimodal_gen.runtime.post_training.rollout_scheduler import (
     prepare_rollout_request_scheduler,
@@ -1526,6 +1527,8 @@ def _denoise_once(
             guidance=None,
         )
         transition = rollout_transition(batch)
+        # UniPC serving does not accept the Euler rollout transition keyword.
+        extra_step_kwargs = {"transition": transition} if transition is not None else {}
 
         do_cfg = guidance_scale > 1.0
         # Control-CFG runs even when text guidance is off (its own extra
@@ -1791,8 +1794,8 @@ def _denoise_once(
                 t,
                 latents,
                 generator=generator,
-                transition=transition,
                 return_dict=False,
+                **extra_step_kwargs,
             )[0]
 
             if action_noise_pred is not None:
@@ -1846,7 +1849,7 @@ def _denoise_once(
             batch=batch,
             latents=latents,
             num_inference_steps=len(timesteps),
-            final_timestep=timesteps.new_zeros(()).cpu(),
+            final_timestep=timesteps.new_zeros(()),
             pipeline_config=server_args.pipeline_config,
         )
 
@@ -1859,14 +1862,17 @@ def _denoise_once(
         return batch
 
     def forward(self, batch: Req, server_args: ServerArgs) -> Req:
-        if batch.extra.get("transfer_plan") is not None:
-            return self._forward_transfer(batch, server_args)
-        with self.use_declared_component(
-            component_name="transformer",
-            module=self.transformer,
-            phase="denoise",
-        ):
-            return self._denoise_once(batch, server_args)
+        try:
+            if batch.extra.get("transfer_plan") is not None:
+                return self._forward_transfer(batch, server_args)
+            with self.use_declared_component(
+                component_name="transformer",
+                module=self.transformer,
+                phase="denoise",
+            ):
+                return self._denoise_once(batch, server_args)
+        finally:
+            release_rollout_resources(batch)
 
     def _predict_noise_cfg(
         self,

python/sglang/multimodal_gen/runtime/post_training/flow_rollout_controller.py

@@ -410,7 +410,7 @@ def _maybe_append_dit_trajectory_step(
         if not self.return_dit_trajectory:
             return
 
-        self.step_timesteps.append(timestep_value.detach().cpu())
+        self.step_timesteps.append(timestep_value.detach())
         return_step_indices = self.return_step_indices
         if return_step_indices is not None and step_index not in return_step_indices:
             return

python/sglang/multimodal_gen/runtime/post_training/rollout_hooks.py

@@ -31,6 +31,12 @@
 )
 
 
+def release_rollout_resources(batch: Req) -> None:
+    """Drop request-owned rollout state, including callbacks capturing the request."""
+    batch.rollout_controllers.clear()
+    batch._rollout_denoising_env_state = None
+
+
 def _kwargs_to_cpu(d: Any) -> Any:
     if isinstance(d, torch.Tensor):
         return d.detach().cpu()
@@ -90,6 +96,7 @@ def prepare_rollout_if_enabled(
     ``build_controller(batch=, pipeline_config=)`` returns it; a loop without a
     scheduler passes its own.
     """
+    release_rollout_resources(batch)
     if not batch.rollout:
         return
     controller = build_controller(batch=batch, pipeline_config=pipeline_config)

python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py

@@ -4,15 +4,23 @@
 
 import types
 import unittest
+from contextlib import nullcontext
+from unittest import mock
 
 import torch
 
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
 from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import (
     FlowMatchEulerDiscreteScheduler,
 )
 from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_unipc_multistep import (
     FlowUniPCMultistepScheduler,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages import (
+    cosmos3,
+)
 from sglang.multimodal_gen.runtime.post_training.rollout_scheduler import (
     prepare_rollout_request_scheduler,
 )
@@ -105,6 +113,62 @@ def test_rl_capable_serving_scheduler_passes_through(self):
         torch.testing.assert_close(serving.sigmas, expected)
 
 
+class TestRolloutRegressions(unittest.TestCase):
+    def test_cosmos3_serving_with_unipc(self):
+        stage = cosmos3.Cosmos3DenoisingStage.__new__(cosmos3.Cosmos3DenoisingStage)
+        stage.scheduler = _serving_scheduler()
+        stage.transformer = mock.Mock()
+        stage._logged_parallel_config = True
+        stage.log_info = mock.Mock()
+        stage.progress_bar = lambda iterable, **kwargs: iterable
+        stage.use_declared_component = lambda **kwargs: nullcontext()
+        stage._run_transformer = lambda **kwargs: torch.zeros_like(kwargs["latents"])
+        text = torch.ones(1, 2, dtype=torch.long)
+        batch = Req(
+            sampling_params=SamplingParams(prompt="test", guidance_scale=1.0),
+            latents=torch.ones(1, 1, 1, 1, 1),
+            timesteps=stage.scheduler.timesteps,
+            extra={
+                "cond_text_ids": text,
+                "cond_text_mask": text,
+                "uncond_text_ids": text,
+                "uncond_text_mask": text,
+                "cond_text_seq_len": 2,
+                "video_shape": (1, 1, 1),
+            },
+        )
+        args = types.SimpleNamespace(pipeline_config=None, enable_cfg_parallel=False)
+        with mock.patch.object(cosmos3, "get_sp_world_size", return_value=1):
+            stage.forward(batch, args)
+        torch.testing.assert_close(batch.latents, torch.ones_like(batch.latents))
+
+    def test_denoiser_failure_releases_rollout_state(self):
+        for stage_cls in (DenoisingStage, cosmos3.Cosmos3DenoisingStage):
+            with self.subTest(stage=stage_cls.__name__):
+                batch = Req(sampling_params=SamplingParams(prompt="test", rollout=True))
+                batch.rollout_controllers["visual"] = object()
+                batch._rollout_denoising_env_state = {"conditioning": torch.zeros(1)}
+                stage = stage_cls.__new__(stage_cls)
+                stage.transformer = None
+                stage.use_declared_component = lambda **kwargs: nullcontext()
+                stage._offload_for_torch_compile_warmup = lambda batch: nullcontext()
+                stage._denoise = stage._denoise_once = mock.Mock(
+                    side_effect=RuntimeError("denoising failed")
+                )
+                with self.assertRaisesRegex(RuntimeError, "denoising failed"):
+                    stage.forward(batch, None)
+                self.assertEqual(batch.rollout_controllers, {})
+                self.assertIsNone(batch._rollout_denoising_env_state)
+
+    def test_request_logging_does_not_copy_controller_tensors(self):
+        batch = Req(sampling_params=SamplingParams(prompt="test"))
+        batch.rollout_controllers["visual"] = {"noise_buffer": torch.zeros(1)}
+        with mock.patch.object(
+            torch.Tensor, "__deepcopy__", side_effect=AssertionError
+        ):
+            self.assertIn("'rollout_controllers': ['visual']", str(batch))
+
+
 class _FusedParamModule(torch.nn.Module):
     """Diffusers-style q/k/v weights that map into one fused to_qkv param."""
 

python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py

@@ -77,6 +77,15 @@ def _build_batch(self, *, debug_mode: bool) -> types.SimpleNamespace:
             latents=torch.zeros(2, 4, 8, 8, dtype=torch.float32),
         )
 
+    @unittest.skipUnless(torch.cuda.is_available(), "CUDA required")
+    def test_trajectory_timestep_stays_on_device(self):
+        batch = self._build_batch(debug_mode=False)
+        batch.rollout_return_dit_trajectory = True
+        rollout = _make_rollout(batch)
+        timestep = torch.tensor(500.0, device="cuda")
+        rollout._maybe_append_dit_trajectory_step(batch.latents, timestep, 0)
+        self.assertEqual(rollout.step_timesteps[0].device, timestep.device)
+
     def test_ode_step_does_not_call_variance_noise_sampler(self):
         batch = self._build_batch(debug_mode=False)
         rollout = _make_rollout(batch)