[Refactor] diffusion: Rollout controller as a scheduler-step plugin instead of mixins

python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py

@@ -134,6 +134,9 @@ def _reject_non_lora_delta_tensors(adapter: dict[str, torch.Tensor]) -> None:
         )
 
 
+# DO NOT CHANGE, MILES-D IS USING: miles_diffusion keeps a clone in
+# miles/backends/fsdp_utils/models/diffusers/h3/weight_bridge.py; change both together.
+# TODO(rockdu): load weights and update weights in one format so the clone can go.
 def _diffusers_h3_checkpoint(
     iterator: Iterable[tuple[str, torch.Tensor]],
 ) -> Iterator[tuple[str, torch.Tensor]]:

python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py

@@ -22,7 +22,7 @@
 # ==============================================================================
 import math
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Callable
 
 import numpy as np
 import torch
@@ -35,9 +35,6 @@
 from sglang.multimodal_gen.runtime.models.schedulers.sigma_schedule import (
     SigmaScheduleMixin,
 )
-from sglang.multimodal_gen.runtime.post_training.scheduler_rl_mixin import (
-    SchedulerRLMixin,
-)
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
@@ -58,7 +55,7 @@ class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
 
 
 class FlowMatchEulerDiscreteScheduler(
-    SchedulerMixin, ConfigMixin, BaseScheduler, SchedulerRLMixin, SigmaScheduleMixin
+    SchedulerMixin, ConfigMixin, BaseScheduler, SigmaScheduleMixin
 ):
     """
     Euler scheduler.
@@ -465,7 +462,7 @@ def step(
         s_noise: float = 1.0,
         generator: torch.Generator | list[torch.Generator] | None = None,
         per_token_timesteps: torch.Tensor | None = None,
-        batch=None,
+        transition: Callable[..., torch.FloatTensor] | None = None,
         return_dict: bool = True,
     ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.FloatTensor, ...]:
         """
@@ -488,6 +485,10 @@ def step(
                 A random number generator, or one generator per batch item.
             per_token_timesteps (`torch.Tensor`, *optional*):
                 The timesteps for each token in the sample.
+            transition (`Callable`, *optional*):
+                Replaces this scheduler's own `_transition` for this call, with the
+                same keyword arguments; rollout passes
+                `FlowRolloutController.transition` here.
             return_dict (`bool`):
                 Whether or not to return a
                 [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple.
@@ -535,24 +536,18 @@ def step(
             next_sigma = sigma_next
             dt = sigma_next - sigma
 
-        if batch is not None and batch.rollout:
-            if not self.already_prepared_rollout(batch):
-                raise RuntimeError("Rollout not prepared before step")
-            prev_sample = self.flow_sde_sampling(
-                batch, model_output, sample, current_sigma, next_sigma, generator
-            )
-        else:
-            if self.config.stochastic_sampling:
-                x0 = sample - current_sigma * model_output
-                noise = randn_tensor(
-                    sample.shape,
-                    generator=generator,
-                    device=sample.device,
-                    dtype=sample.dtype,
-                )
-                prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise
-            else:
-                prev_sample = sample + dt * model_output
+        if transition is None:
+            transition = self._transition
+        prev_sample = transition(
+            model_output=model_output,
+            sample=sample,
+            current_sigma=current_sigma,
+            next_sigma=next_sigma,
+            dt=dt,
+            generator=generator,
+            step_index=self.step_index,
+            timestep=timestep,
+        )
 
         # upon completion increase step index by one
         assert self._step_index is not None, "_step_index should not be None"
@@ -566,6 +561,33 @@ def step(
 
         return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
 
+    # Old ODE transition: the scheduler's own Euler update (or its stochastic_sampling step),
+    # used when step() is not given a transition.
+    def _transition(
+        self,
+        *,
+        model_output: torch.FloatTensor,
+        sample: torch.FloatTensor,
+        current_sigma: torch.FloatTensor,
+        next_sigma: torch.FloatTensor,
+        dt: torch.FloatTensor,
+        generator: torch.Generator | list[torch.Generator] | None,
+        step_index: int,
+        timestep: int | torch.Tensor,
+    ) -> torch.FloatTensor:
+        if self.config.stochastic_sampling:
+            x0 = sample - current_sigma * model_output
+            noise = randn_tensor(
+                sample.shape,
+                generator=generator,
+                device=sample.device,
+                dtype=sample.dtype,
+            )
+            prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise
+        else:
+            prev_sample = sample + dt * model_output
+        return prev_sample
+
     def _time_shift_exponential(
         self, mu: float, sigma: float, t: torch.Tensor | np.ndarray
     ) -> torch.Tensor | np.ndarray:

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

@@ -27,6 +27,9 @@
     DataType,
     SamplingParams,
 )
+from sglang.multimodal_gen.runtime.post_training.flow_rollout_controller import (
+    FlowRolloutController,
+)
 from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import (
     RolloutTrajectoryData,
 )
@@ -198,6 +201,8 @@ class Req:
     trajectory_timesteps: torch.Tensor | None = None
     trajectory_latents: torch.Tensor | None = None
     rollout_trajectory_data: RolloutTrajectoryData | None = None
+    # Rollout controllers of this request, keyed by latent stream.
+    rollout_controllers: dict[str, FlowRolloutController] = field(default_factory=dict)
     trajectory_audio_latents: torch.Tensor | None = None
 
     # Extra parameters that might be needed by specific pipeline implementations (e.g., LTX2.3 DenoisingAVStage)

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

@@ -141,8 +141,11 @@
     AttentionBackendEnum,
     current_platform,
 )
-from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import (
-    RolloutDenoisingMixin,
+from sglang.multimodal_gen.runtime.post_training.rollout_hooks import (
+    finalize_rollout_if_enabled,
+    init_rollout_env_if_enabled,
+    prepare_rollout_if_enabled,
+    rollout_step_kwargs,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -290,7 +293,7 @@ class DualTransformerExecutionMode(str, Enum):
     PAIRED_PER_STEP = "paired_per_step"
 
 
-class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
+class DenoisingStage(PipelineStage):
     """
     Stage for running the denoising loop in diffusion pipelines.
 
@@ -1251,13 +1254,18 @@ def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs):
         if freshly_loaded:
             register_loaded_transformer(self, server_args, pipeline)
 
-        if batch.rollout:
-            self._maybe_prepare_rollout(batch)
+        prepare_rollout_if_enabled(
+            batch=batch, pipeline_config=server_args.pipeline_config
+        )
 
         # Prepare extra step kwargs for scheduler
         extra_step_kwargs = self.prepare_extra_func_kwargs(
             scheduler.step,
-            {"generator": batch.generator, "eta": batch.eta, "batch": batch},
+            {
+                "generator": batch.generator,
+                "eta": batch.eta,
+                **rollout_step_kwargs(batch),
+            },
         )
 
         # Setup precision and autocast settings
@@ -2027,26 +2035,23 @@ def _denoise(
         Run the denoising loop.
         """
         ctx = self._prepare_denoising_loop(batch, server_args)
-        if batch.rollout:
-            self._maybe_init_denoising_env_collection(
-                batch=batch,
-                pipeline_config=server_args.pipeline_config,
-                image_kwargs=ctx.image_kwargs,
-                pos_cond_kwargs=ctx.pos_cond_kwargs,
-                neg_cond_kwargs=ctx.neg_cond_kwargs,
-                guidance=ctx.guidance,
-            )
+        init_rollout_env_if_enabled(
+            batch=batch,
+            image_kwargs=ctx.image_kwargs,
+            pos_cond_kwargs=ctx.pos_cond_kwargs,
+            neg_cond_kwargs=ctx.neg_cond_kwargs,
+            guidance=ctx.guidance,
+        )
         timesteps_cpu = self._run_denoising_loop(ctx, batch, server_args)
 
         # collect rollout outputs before finalization gathers or replaces latents
-        if batch.rollout:
-            self._postprocess_rollout_outputs(
-                batch=batch,
-                latents=ctx.latents,
-                num_inference_steps=len(timesteps_cpu),
-                final_timestep=timesteps_cpu.new_zeros(()),
-                server_args=server_args,
-            )
+        finalize_rollout_if_enabled(
+            batch=batch,
+            latents=ctx.latents,
+            num_inference_steps=len(timesteps_cpu),
+            final_timestep=ctx.timesteps.new_zeros(()),
+            pipeline_config=server_args.pipeline_config,
+        )
         self._finalize_denoising_loop(ctx, batch, server_args)
         return batch
 
@@ -2108,19 +2113,6 @@ def _run_denoising_loop(
                             t_host,
                             timesteps_cpu,
                         )
-                        # Capture the raw (pre-scale, pre-I2V-concat) noisy latent
-                        # x_{t_i} for rollout trajectory collection. Must run
-                        # BEFORE _run_denoising_step so ctx.latents is still the
-                        # pre-step value. Gated on batch.rollout to keep the
-                        # non-rollout path strictly untouched.
-                        if collect_trajectory and batch.rollout:
-                            batch._rollout_loop_step_index = step_index
-                            self._maybe_append_dit_trajectory_step(
-                                batch=batch,
-                                latents=ctx.latents,
-                                timestep_value=step.t_host,
-                                step_index=step_index,
-                            )
                         self._run_denoising_step(ctx, step, batch, server_args)
                         if collect_trajectory:
                             self._record_trajectory(ctx, step, batch, server_args)

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

@@ -59,8 +59,11 @@
     VerificationResult,
 )
 from sglang.multimodal_gen.runtime.platforms import current_platform
-from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import (
-    RolloutDenoisingMixin,
+from sglang.multimodal_gen.runtime.post_training.rollout_hooks import (
+    finalize_rollout_if_enabled,
+    init_rollout_env_if_enabled,
+    prepare_rollout_if_enabled,
+    rollout_step_kwargs,
 )
 from sglang.multimodal_gen.runtime.post_training.rollout_scheduler import (
     prepare_rollout_request_scheduler,
@@ -1059,7 +1062,7 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         return batch
 
 
-class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin):
+class Cosmos3DenoisingStage(PipelineStage):
     """Cosmos3 denoise loop, including CFG and the parallelism modes.
 
     The UND pathway runs once and its K/V is cached per cache_key (``cond`` /
@@ -1502,23 +1505,24 @@ def _denoise_once(
                 raise ValueError(
                     "Cosmos3 rollout does not support action/sound modalities."
                 )
-            self._maybe_prepare_rollout(batch)
-            self._maybe_init_denoising_env_collection(
-                batch=batch,
-                pipeline_config=server_args.pipeline_config,
-                image_kwargs={},
-                pos_cond_kwargs={
-                    "text_ids": cond_text_ids,
-                    "text_mask": cond_text_mask,
-                    "fps": fps,
-                },
-                neg_cond_kwargs={
-                    "text_ids": uncond_text_ids,
-                    "text_mask": uncond_text_mask,
-                    "fps": fps,
-                },
-                guidance=None,
-            )
+        prepare_rollout_if_enabled(
+            batch=batch, pipeline_config=server_args.pipeline_config
+        )
+        init_rollout_env_if_enabled(
+            batch=batch,
+            image_kwargs={},
+            pos_cond_kwargs={
+                "text_ids": cond_text_ids,
+                "text_mask": cond_text_mask,
+                "fps": fps,
+            },
+            neg_cond_kwargs={
+                "text_ids": uncond_text_ids,
+                "text_mask": uncond_text_mask,
+                "fps": fps,
+            },
+            guidance=None,
+        )
 
         do_cfg = guidance_scale > 1.0
         # Control-CFG runs even when text guidance is off (its own extra
@@ -1779,31 +1783,14 @@ def _denoise_once(
             if velocity_mask is not None:
                 noise_pred = noise_pred * velocity_mask
 
-            if batch.rollout:
-                # Capture the pre-step x_{t_i} before the scheduler advances it.
-                batch._rollout_loop_step_index = i
-                self._maybe_append_dit_trajectory_step(
-                    batch=batch,
-                    latents=latents,
-                    timestep_value=t,
-                    step_index=i,
-                )
-                latents = scheduler.step(
-                    noise_pred,
-                    t,
-                    latents,
-                    generator=batch.generator,
-                    batch=batch,
-                    return_dict=False,
-                )[0]
-            else:
-                latents = scheduler.step(
-                    noise_pred,
-                    t,
-                    latents,
-                    generator=generator,
-                    return_dict=False,
-                )[0]
+            latents = scheduler.step(
+                noise_pred,
+                t,
+                latents,
+                generator=generator,
+                **rollout_step_kwargs(batch),
+                return_dict=False,
+            )[0]
 
             if action_noise_pred is not None:
                 # Zero the velocity at conditioned (clean) action tokens and at
@@ -1852,14 +1839,13 @@ def _denoise_once(
         # next request's steps.
         self.transformer.reset_denoising_step()
 
-        if batch.rollout:
-            self._postprocess_rollout_outputs(
-                batch=batch,
-                latents=latents,
-                num_inference_steps=len(timesteps),
-                final_timestep=timesteps.new_zeros(()).cpu(),
-                server_args=server_args,
-            )
+        finalize_rollout_if_enabled(
+            batch=batch,
+            latents=latents,
+            num_inference_steps=len(timesteps),
+            final_timestep=timesteps.new_zeros(()),
+            pipeline_config=server_args.pipeline_config,
+        )
 
         batch.latents = latents
         if action_latents is not None: