[diffusion] Support RL rollout for the Wan pipeline via a per-request scheduler switch

python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py

@@ -8,6 +8,9 @@
 using the modular pipeline architecture.
 """
 
+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,
 )
@@ -17,6 +20,7 @@
 from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline
 from sglang.multimodal_gen.runtime.pipelines_core.stages import (
     InputValidationStage,
+    TimestepPreparationStage,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.wan import (
     WanProgressiveDenoisingStage,
@@ -49,10 +53,20 @@ def initialize_pipeline(self, server_args: ServerArgs):
         )
 
     def create_pipeline_stages(self, server_args: ServerArgs) -> None:
+        shift = server_args.pipeline_config.flow_shift
         self.add_stage(InputValidationStage())
         self.add_standard_text_encoding_stage()
         self.add_standard_latent_preparation_stage()
-        self.add_standard_timestep_preparation_stage()
+        # Serving keeps UniPC; requests with rollout=True bind a first-order
+        # flow-match Euler scheduler, which the RL SDE/log-prob path requires.
+        self.add_stage(
+            TimestepPreparationStage(
+                scheduler=self.get_module("scheduler"),
+                rollout_scheduler=FlowMatchEulerDiscreteScheduler(
+                    shift=1.0 if shift is None else shift
+                ),
+            )
+        )
         self.add_progressive_denoising_stage(WanProgressiveDenoisingStage)
         self.add_standard_decoding_stage()
 

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

@@ -28,6 +28,9 @@
 from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
     VerificationResult,
 )
+from sglang.multimodal_gen.runtime.post_training.rollout_timestep_mixin import (
+    RolloutTimestepPreparationMixin,
+)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
@@ -45,7 +48,7 @@ class TimestepPreparationFingerprint:
     num_frames: int | None
 
 
-class TimestepPreparationStage(PipelineStage):
+class TimestepPreparationStage(PipelineStage, RolloutTimestepPreparationMixin):
     """
     Stage for preparing timesteps for the diffusion process.
 
@@ -63,9 +66,12 @@ def __init__(
         prepare_extra_set_timesteps_kwargs: (
             list[Callable[[Req, ServerArgs], Tuple[str, Any]]] | None
         ) = None,
+        rollout_scheduler=None,
     ) -> None:
         super().__init__()
         self.scheduler = scheduler
+        # See RolloutTimestepPreparationMixin.
+        self.rollout_scheduler = rollout_scheduler
         self.prepare_extra_set_timesteps_kwargs = list(
             prepare_extra_set_timesteps_kwargs or []
         )
@@ -90,7 +96,11 @@ def forward(
         if batch.scheduler is not None and batch.timesteps is not None:
             return batch
 
-        scheduler = get_or_create_request_scheduler(batch, self.scheduler)
+        rollout_template = self._resolve_rollout_scheduler(batch)
+        scheduler = get_or_create_request_scheduler(
+            batch,
+            rollout_template if rollout_template is not None else self.scheduler,
+        )
         device = get_local_torch_device()
         num_inference_steps = batch.num_inference_steps
         timesteps = batch.timesteps
@@ -153,6 +163,9 @@ def forward(
             )
             timesteps = scheduler.timesteps
 
+        if rollout_template is not None:
+            self._check_rollout_timesteps(scheduler)
+
         # Update batch with prepared timesteps
         batch.timesteps = timesteps
         batch.scheduler = scheduler

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

@@ -38,41 +38,52 @@ def _kwargs_to_cpu(d: Any) -> Any:
 
 class RolloutDenoisingMixin:
 
+    def _request_scheduler(self, batch: Req):
+        """Scheduler in effect for this request.
+
+        The timestep preparation stage binds per-request schedulers (e.g. the
+        RL rollout scheduler) to ``batch.scheduler``; the stage module is only
+        a fallback for pipelines that never bind one.
+        """
+        return batch.scheduler if batch.scheduler is not None else self.scheduler
+
     def _maybe_prepare_rollout(self, batch: Req):
         """Prepare denoising loop for rollout."""
-        if not isinstance(self.scheduler, SchedulerRLMixin):
+        scheduler = self._request_scheduler(batch)
+        if not isinstance(scheduler, SchedulerRLMixin):
             if batch.rollout:
                 raise ValueError(
-                    f"Scheduler {type(self.scheduler)} does not support rollout"
+                    f"Scheduler {type(scheduler)} does not support rollout"
                 )
             return
 
-        self.scheduler.release_rollout_resources(batch)
+        scheduler.release_rollout_resources(batch)
         if batch.rollout:
-            self.scheduler.prepare_rollout(
+            scheduler.prepare_rollout(
                 batch=batch,
                 pipeline_config=self.server_args.pipeline_config,
             )
 
     def _maybe_collect_rollout_log_probs(self, batch: Req):
-        if not isinstance(self.scheduler, SchedulerRLMixin):
+        scheduler = self._request_scheduler(batch)
+        if not isinstance(scheduler, SchedulerRLMixin):
             if batch.rollout:
                 raise ValueError(
-                    f"Scheduler {type(self.scheduler)} does not support rollout"
+                    f"Scheduler {type(scheduler)} does not support rollout"
                 )
             return
 
         if batch.rollout:
             if batch.rollout_trajectory_data is None:
                 batch.rollout_trajectory_data = RolloutTrajectoryData()
             batch.rollout_trajectory_data.rollout_log_probs = (
-                self.scheduler.collect_rollout_log_probs(batch)
+                scheduler.collect_rollout_log_probs(batch)
             )
             if batch.rollout_debug_mode:
                 batch.rollout_trajectory_data.rollout_debug_tensors = (
-                    self.scheduler.collect_rollout_debug_tensors(batch)
+                    scheduler.collect_rollout_debug_tensors(batch)
                 )
-            self.scheduler.release_rollout_resources(batch)
+            scheduler.release_rollout_resources(batch)
 
     def _postprocess_rollout_outputs(
         self,

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

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Mixin for per-request rollout scheduler binding in TimestepPreparationStage.
+
+Kept under post_training to keep the core stage lean; mirrors
+RolloutDenoisingMixin on DenoisingStage.
+"""
+
+from __future__ import annotations
+
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class RolloutTimestepPreparationMixin:
+    """Bind an alternate scheduler to rollout=True requests.
+
+    The rollout SDE/log-prob path needs a first-order flow-match Euler
+    scheduler, which not every pipeline serves (e.g. Wan serves UniPC). The
+    host stage sets ``self.rollout_scheduler``; None keeps the serving
+    scheduler for rollout requests. Downstream stages read the scheduler
+    from ``batch.scheduler``, so the host stage is the single switch point.
+    """
+
+    # Class-level so the rollout info log prints once per process, not once
+    # per stage instance.
+    _logged_rollout_scheduler_check = False
+
+    def _resolve_rollout_scheduler(self, batch: Req):
+        """Return the rollout scheduler template for this request, or None."""
+        if batch.rollout and self.rollout_scheduler is not None:
+            return self.rollout_scheduler
+        return None
+
+    def _check_rollout_timesteps(self, scheduler) -> None:
+        # The rollout SDE/log-prob math assumes the flow-match Euler
+        # convention timesteps == sigmas[:-1] * num_train_timesteps.
+        sigmas = scheduler.sigmas
+        timesteps = scheduler.timesteps
+        if sigmas is None or timesteps is None or sigmas.numel() < 2:
+            return
+        reconstructed = sigmas[:-1].to(device=timesteps.device) * float(
+            scheduler.config.num_train_timesteps
+        )
+        max_abs_diff = (timesteps.float() - reconstructed.float()).abs().max().item()
+        if max_abs_diff > 1e-3:
+            raise ValueError(
+                f"rollout timestep/sigma mismatch: max_abs_diff={max_abs_diff:.6g}"
+            )
+        if not RolloutTimestepPreparationMixin._logged_rollout_scheduler_check:
+            logger.info(
+                "RL rollout using %s (timesteps dtype=%s, sigmas dtype=%s, "
+                "max_abs_diff=%.6g)",
+                type(scheduler).__name__,
+                timesteps.dtype,
+                sigmas.dtype,
+                max_abs_diff,
+            )
+            RolloutTimestepPreparationMixin._logged_rollout_scheduler_check = True