[diffusion] RL rollout support for the Cosmos3 pipeline

python/sglang/jit_kernel/flash_attention_v3.py

@@ -17,6 +17,10 @@
 
 
 def _call_fa3_kernel(kernel, *args, out=None, **kwargs):
+    # only_qv=False is the kernel default; drop it so older kernel builds
+    # without the argument keep working.
+    if kwargs.get("only_qv") is False:
+        del kwargs["only_qv"]
     if out is None:
         return kernel(*args, **kwargs)
     try:

python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py

@@ -130,6 +130,7 @@ def _slice_rollout_trajectory_for_sample(
         dit_trajectory = RolloutDitTrajectory(
             latents=_extract_single_sample_tensor(dit.latents, sample_idx, batch_size),
             timesteps=dit.timesteps,
+            sigmas=dit.sigmas,
         )
     return RolloutTrajectoryData(
         rollout_log_probs=log_probs,
@@ -143,6 +144,7 @@ def _serialize_rollout_trajectory(
     rtd: RolloutTrajectoryData | None,
     *,
     serialized_dit_timesteps: dict | None = None,
+    serialized_dit_sigmas: dict | None = None,
 ) -> tuple[dict | None, dict | None, dict | None, dict | None]:
     """Return order: rollout_log_probs, rollout_debug_tensors, denoising_env, dit_trajectory."""
     if rtd is None:
@@ -182,6 +184,7 @@ def _serialize_rollout_trajectory(
                 _maybe_serialize(dit.latents) if dit.latents is not None else None
             ),
             "timesteps": serialized_dit_timesteps,
+            "sigmas": serialized_dit_sigmas,
         }
     return (
         serialized_log_probs,
@@ -211,10 +214,14 @@ def _build_response(
         ), "rollout_trajectory_data must be present when rollout=True"
 
     serialized_dit_timesteps = None
+    serialized_dit_sigmas = None
     if rollout and rollout_trajectory_data and rollout_trajectory_data.dit_trajectory:
         serialized_dit_timesteps = _maybe_serialize(
             rollout_trajectory_data.dit_trajectory.timesteps
         )
+        serialized_dit_sigmas = _maybe_serialize(
+            rollout_trajectory_data.dit_trajectory.sigmas
+        )
 
     responses: list[RolloutResponse] = []
     for sample_idx in range(batch_size):
@@ -243,6 +250,7 @@ def _build_response(
         ) = _serialize_rollout_trajectory(
             per_sample_trajectory,
             serialized_dit_timesteps=serialized_dit_timesteps,
+            serialized_dit_sigmas=serialized_dit_sigmas,
         )
         responses.append(
             RolloutResponse(

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

@@ -8,6 +8,9 @@
 
 import os
 
+from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )
@@ -84,7 +87,14 @@ def create_pipeline_stages(self, server_args: ServerArgs) -> None:
 
             self.add_stage(Cosmos3TextGuardrailStage())
         self.add_stage(Cosmos3LatentPreparationStage(vae, transformer))
-        self.add_stage(Cosmos3TimestepPreparationStage(scheduler))
+        # rollout=True requests lazily bind a flow-match Euler scheduler (RL SDE
+        # path); it inherits the serving scheduler's sigma grid per request.
+        self.add_stage(
+            Cosmos3TimestepPreparationStage(
+                scheduler,
+                rollout_scheduler_factory=FlowMatchEulerDiscreteScheduler,
+            )
+        )
         self.add_stage(Cosmos3DenoisingStage(transformer, scheduler, server_args))
         self.add_stage(Cosmos3DecodingStage(vae, guardrails=guardrails_on))
 

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

@@ -27,6 +27,9 @@
     get_sp_world_size,
 )
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import (
+    get_or_create_request_scheduler,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
     PipelineStage,
@@ -39,6 +42,12 @@
     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_timestep_mixin import (
+    RolloutTimestepPreparationMixin,
+)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler
@@ -313,6 +322,8 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         generator = batch.generator
         if generator is None and batch.seed is not None:
             generator = torch.Generator(device=device).manual_seed(batch.seed)
+            # The rollout SDE step draws its variance noise from this generator.
+            batch.generator = generator
 
         noise = torch.randn(shape, generator=generator, device=device, dtype=dtype)
 
@@ -350,7 +361,7 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         return batch
 
 
-class Cosmos3TimestepPreparationStage(PipelineStage):
+class Cosmos3TimestepPreparationStage(PipelineStage, RolloutTimestepPreparationMixin):
     """
     Timestep preparation stage for Cosmos3.
 
@@ -359,12 +370,15 @@ class Cosmos3TimestepPreparationStage(PipelineStage):
 
     parallelism_type = StageParallelismType.REPLICATED
 
-    def __init__(self, scheduler):
+    def __init__(self, scheduler, rollout_scheduler_factory=None):
         super().__init__()
         self.scheduler = scheduler
         self.default_flow_shift = getattr(
             getattr(scheduler, "config", None), "flow_shift", None
         )
+        # See RolloutTimestepPreparationMixin.
+        self.rollout_scheduler_factory = rollout_scheduler_factory
+        self._rollout_scheduler = None
 
     def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         """Prepare scheduler timesteps."""
@@ -381,13 +395,34 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         self.scheduler.set_timesteps(num_inference_steps, device=device)
         batch.timesteps = self.scheduler.timesteps
 
+        rollout_template = self._resolve_rollout_scheduler(batch)
+        if rollout_template is not None:
+            scheduler = get_or_create_request_scheduler(batch, rollout_template)
+            explicit_shift = getattr(batch, "flow_shift", None)
+            if explicit_shift is None:
+                explicit_shift = server_args.pipeline_config.flow_shift
+            if explicit_shift is not None:
+                # An explicit flow_shift selects a plain shifted grid — the
+                # checkpoint's karras schedule ignores flow_shift entirely,
+                # and its dense head starves the RL gradient (dt ~ 1e-3).
+                scheduler.set_shift(float(explicit_shift))
+                scheduler.set_timesteps(num_inference_steps, device=device)
+            else:
+                # Reuse the serving scheduler's sigma grid (sans terminal
+                # sigma) so rollout noise levels match serving exactly.
+                scheduler.set_timesteps(
+                    sigmas=self.scheduler.sigmas[:-1].tolist(), device=device
+                )
+            batch.timesteps = scheduler.timesteps
+            self._check_rollout_timesteps(scheduler)
+
         self.log_info(
             f"Prepared {len(batch.timesteps)} timesteps (flow_shift={flow_shift})"
         )
         return batch
 
 
-class Cosmos3DenoisingStage(PipelineStage):
+class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin):
     """Cosmos3 denoise loop, including CFG and the parallelism modes.
 
     The UND pathway runs once and its K/V is cached per cache_key (``cond`` /
@@ -578,6 +613,31 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         image_latent = batch.image_latent
         guidance_interval = getattr(batch.sampling_params, "guidance_interval", None)
 
+        scheduler = batch.scheduler if batch.scheduler is not None else self.scheduler
+        if batch.rollout:
+            if image_latent is not None:
+                raise ValueError(
+                    "Cosmos3 rollout supports T2V/T2I only; I2V frame-0 "
+                    "re-injection breaks the Gaussian transition assumption."
+                )
+            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,
+            )
+
         do_cfg = guidance_scale > 1.0
 
         enable_cfg_parallel = server_args.enable_cfg_parallel and do_cfg
@@ -696,19 +756,45 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req:
             if velocity_mask is not None:
                 noise_pred = noise_pred * velocity_mask
 
-            latents = self.scheduler.step(
-                noise_pred,
-                t,
-                latents,
-                return_dict=False,
-            )[0]
+            if batch.rollout:
+                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,
+                    return_dict=False,
+                )[0]
 
             if image_latent is not None:
                 latents[:, :, 0:1, :, :] = image_latent
 
             if batch.profile and not batch.is_warmup:
                 self.step_profile()
 
+        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,
+            )
+
         batch.latents = latents
         self.log_info("Denoising complete")
         return batch
@@ -938,4 +1024,5 @@ def forward(self, batch: Req, server_args: ServerArgs):
         return OutputBatch(
             output=output,
             metrics=batch.metrics if hasattr(batch, "metrics") else None,
+            rollout_trajectory_data=batch.rollout_trajectory_data,
         )

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

@@ -52,6 +52,10 @@ class RolloutDitTrajectory:
     # final denoised latent x_{t_T} (last scheduler.step output).
     latents: torch.Tensor | None = None
     timesteps: torch.Tensor | None = None  # [T]
+    # Scheduler sigma grid [T+1] including the terminal sigma. Lets the
+    # training side replay the exact rollout noise levels instead of
+    # reconstructing them from timesteps (which drifts in fp32).
+    sigmas: torch.Tensor | None = None
 
 
 @dataclass

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,
@@ -181,9 +192,15 @@ def _maybe_finalize_denoising_env_collection(self, batch, pipeline_config) -> No
                 batch=batch,
                 stacked_latents=step_latents_tensor,
             )
+            scheduler_sigmas = self._request_scheduler(batch).sigmas
             batch.rollout_trajectory_data.dit_trajectory = RolloutDitTrajectory(
                 latents=step_latents_tensor.cpu(),
                 timesteps=torch.stack(step_timesteps, dim=0).cpu(),
+                sigmas=(
+                    scheduler_sigmas.detach().cpu()
+                    if scheduler_sigmas is not None
+                    else None
+                ),
             )
 
         if env is not None and batch.rollout_return_denoising_env: