[diffusion] rl: support standalone rollout api, denoising environment backpass and sp-aligned log-prob for T2I post-training (#22604)

Co-authored-by: MikukuOvO <mikukuovo@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -508,6 +508,9 @@ class PipelineConfig:
     def _unpad_and_unpack_latents(self, latents, audio_latents, batch, vae, audio_vae):
         raise NotImplementedError("not yet implemented")
 
+    def gather_dit_env_static_for_sp(self, batch, cond_kwargs: dict | None):
+        return cond_kwargs
+
     @staticmethod
     def add_cli_args(
         parser: FlexibleArgumentParser, prefix: str = ""

python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py

@@ -18,6 +18,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     maybe_unpad_latents,
     shard_rotary_emb_for_sp,
 )
+from sglang.multimodal_gen.configs.post_training.pipeline_configs import (
+    QwenImageRolloutPipelineMixin,
+)
 from sglang.multimodal_gen.runtime.models.vision_utils import resize
 from sglang.multimodal_gen.utils import calculate_dimensions
 
@@ -127,7 +130,7 @@ def _pack_latents(latents, batch_size, num_channels_latents, height, width):
 
 
 @dataclass
-class QwenImagePipelineConfig(ImagePipelineConfig):
+class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig):
     """Configuration for the QwenImage pipeline."""
 
     should_use_guidance: bool = False

python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py

@@ -15,6 +15,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ImagePipelineConfig,
     ModelTaskType,
 )
+from sglang.multimodal_gen.configs.post_training.pipeline_configs import (
+    ZImageRolloutPipelineMixin,
+)
 from sglang.multimodal_gen.runtime.distributed.parallel_state import (
     get_sp_group,
     get_sp_parallel_rank,
@@ -40,7 +43,7 @@ class TransformersModelConfig(EncoderConfig):
 
 
 @dataclass
-class ZImagePipelineConfig(ImagePipelineConfig):
+class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
     should_use_guidance: bool = False
     task_type: ModelTaskType = ModelTaskType.T2I
     dit_config: DiTConfig = field(default_factory=ZImageDitConfig)

python/sglang/multimodal_gen/configs/post_training/pipeline_configs/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Rollout / RL hooks mixed into multimodal pipeline configs."""
+
+from sglang.multimodal_gen.configs.post_training.pipeline_configs.qwen_image_rollout_pipeline_mixin import (
+    QwenImageRolloutPipelineMixin,
+)
+from sglang.multimodal_gen.configs.post_training.pipeline_configs.zimage_rollout_pipeline_mixin import (
+    ZImageRolloutPipelineMixin,
+)
+
+__all__ = [
+    "QwenImageRolloutPipelineMixin",
+    "ZImageRolloutPipelineMixin",
+]

python/sglang/multimodal_gen/configs/post_training/pipeline_configs/qwen_image_rollout_pipeline_mixin.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Rollout / RL hooks for Qwen-Image pipeline configs."""
+
+from __future__ import annotations
+
+import torch
+
+from sglang.multimodal_gen.runtime.post_training.sp_utils import (
+    all_gather_if_sp_sharded,
+    maybe_trim_sp_rope_seq_for_batch,
+)
+
+
+class QwenImageRolloutPipelineMixin:
+
+    def gather_dit_env_static_for_sp(self, batch, cond_kwargs: dict | None):
+        if cond_kwargs is None:
+            return None
+        out = dict(cond_kwargs)
+        freqs = out.get("freqs_cis")
+        if freqs is not None:
+            img_cache, txt_cache = freqs[0], freqs[1]
+            if isinstance(img_cache, torch.Tensor) and img_cache.dim() == 2:
+                img_g = all_gather_if_sp_sharded(batch, img_cache, dim=0)
+                img_g = maybe_trim_sp_rope_seq_for_batch(batch, img_g)
+                out["freqs_cis"] = (img_g, txt_cache)
+        return out

python/sglang/multimodal_gen/configs/post_training/pipeline_configs/zimage_rollout_pipeline_mixin.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Rollout / RL hooks for Z-Image pipeline configs."""
+
+from __future__ import annotations
+
+import torch
+
+from sglang.multimodal_gen.runtime.post_training.sp_utils import (
+    all_gather_if_sp_sharded,
+    maybe_trim_sp_rope_seq_for_batch,
+)
+
+
+class ZImageRolloutPipelineMixin:
+
+    def gather_dit_env_static_for_sp(self, batch, cond_kwargs: dict | None):
+        if cond_kwargs is None:
+            return None
+        out = dict(cond_kwargs)
+        freqs = out.get("freqs_cis")
+        if freqs is not None:
+            cap_freqs, x_freqs = freqs[0], freqs[1]
+            if isinstance(x_freqs, torch.Tensor) and x_freqs.dim() >= 2:
+                x_g = all_gather_if_sp_sharded(batch, x_freqs, dim=0)
+                x_g = maybe_trim_sp_rope_seq_for_batch(batch, x_g)
+                out["freqs_cis"] = (cap_freqs, x_g)
+        return out

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -187,6 +187,12 @@ class SamplingParams:
     )
     return_trajectory_latents: bool = False  # returns all latents for each timestep
     return_trajectory_decoded: bool = False  # returns decoded latents for each timestep
+    rollout_return_denoising_env: bool = (
+        False  # populate ``denoising_env`` (image/pos/neg kwargs, guidance) for RL replay
+    )
+    rollout_return_dit_trajectory: bool = (
+        False  # per-step noisy latents + final latent + timesteps (RolloutDitTrajectory)
+    )
     # if True, disallow user params to override subclass-defined protected fields
     no_override_protected_fields: bool = False
     # whether to adjust num_frames for multi-GPU friendly splitting (default: True)

python/sglang/multimodal_gen/runtime/entrypoints/http_server.py

@@ -16,7 +16,10 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
     VertexGenerateReqInput,
 )
 from sglang.multimodal_gen.runtime.entrypoints.openai.utils import build_sampling_params
-from sglang.multimodal_gen.runtime.entrypoints.post_training import weights_api
+from sglang.multimodal_gen.runtime.entrypoints.post_training import (
+    rollout_api,
+    weights_api,
+)
 from sglang.multimodal_gen.runtime.entrypoints.utils import (
     prepare_request,
     save_outputs,
@@ -282,6 +285,7 @@ def create_app(server_args: ServerArgs):
     app.include_router(video_api.router)
     app.include_router(mesh_api.router)
     app.include_router(weights_api.router)
+    app.include_router(rollout_api.router)
 
     app.state.server_args = server_args
     return app

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

@@ -1,6 +1,11 @@
 """Request/response data structures for post-training APIs."""
 
+from __future__ import annotations
+
 from dataclasses import dataclass
+from typing import Any, Optional
+
+from pydantic import BaseModel
 
 
 @dataclass
@@ -17,3 +22,54 @@ class GetWeightsChecksumReqInput:
     """Compute SHA-256 checksum of loaded module weights for verification."""
 
     module_names: list[str] | None = None
+
+
+class RolloutRequest(BaseModel):
+    prompt: str
+    negative_prompt: Optional[str] = None
+    seed: int = 1024
+    generator_device: str = "cuda"
+
+    width: Optional[int] = None
+    height: Optional[int] = None
+    num_inference_steps: Optional[int] = None
+    num_outputs_per_prompt: Optional[int] = None
+
+    guidance_scale: Optional[float] = None
+    true_cfg_scale: Optional[float] = None
+
+    # video-specific (ignored by image pipelines)
+    num_frames: Optional[int] = None
+    fps: Optional[int] = None
+
+    rollout: bool = True
+    rollout_sde_type: str = "sde"
+    rollout_noise_level: float = 0.7
+    rollout_log_prob_no_const: bool = False
+    rollout_debug_mode: bool = True
+
+    rollout_return_denoising_env: bool = False
+    rollout_return_dit_trajectory: bool = False
+
+    image_path: Optional[list[str]] = None
+
+    # suppress verbose per-request logging (also gates peak_memory_mb collection)
+    suppress_logs: bool = False
+
+    extra_sampling_params: Optional[dict[str, Any]] = None
+
+
+class RolloutResponse(BaseModel):
+    request_id: str
+    prompt: str
+    seed: int
+
+    generated_output: Any = None
+
+    rollout_log_probs: Optional[dict[str, Any]] = None
+    rollout_debug_tensors: Optional[dict[str, Any]] = None
+    denoising_env: Optional[dict[str, Any]] = None
+    dit_trajectory: Optional[dict[str, Any]] = None
+
+    inference_time_s: Optional[float] = None
+    peak_memory_mb: Optional[float] = None

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

@@ -0,0 +1,309 @@
+"""Rollout HTTP API (``POST /rollout/generate``)."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from fastapi import APIRouter, HTTPException
+from fastapi.responses import ORJSONResponse
+
+from sglang.multimodal_gen.configs.sample.sampling_params import generate_request_id
+from sglang.multimodal_gen.runtime.entrypoints.openai.utils import build_sampling_params
+from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import (
+    RolloutRequest,
+    RolloutResponse,
+)
+from sglang.multimodal_gen.runtime.entrypoints.post_training.utils import (
+    _maybe_serialize,
+)
+from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
+from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import (
+    RolloutDebugTensors,
+    RolloutDenoisingEnv,
+    RolloutDitTrajectory,
+    RolloutTrajectoryData,
+)
+from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
+from sglang.multimodal_gen.runtime.server_args import get_global_server_args
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+router = APIRouter(prefix="/rollout", tags=["rollout"])
+
+
+def _extract_single_sample_tensor(obj: Any, sample_idx: int, batch_size: int) -> Any:
+    if isinstance(obj, torch.Tensor):
+        if obj.dim() >= 1 and obj.shape[0] == batch_size:
+            return obj[sample_idx].contiguous()
+        return obj
+    if isinstance(obj, dict):
+        return {
+            k: _extract_single_sample_tensor(v, sample_idx, batch_size)
+            for k, v in obj.items()
+        }
+    if isinstance(obj, list):
+        return [_extract_single_sample_tensor(v, sample_idx, batch_size) for v in obj]
+    if isinstance(obj, tuple):
+        return tuple(
+            _extract_single_sample_tensor(v, sample_idx, batch_size) for v in obj
+        )
+    return obj
+
+
+def _slice_rollout_trajectory_for_sample(
+    rtd: RolloutTrajectoryData | None,
+    sample_idx: int,
+    batch_size: int,
+) -> RolloutTrajectoryData | None:
+    if rtd is None:
+        return None
+    log_probs = rtd.rollout_log_probs
+    if (
+        isinstance(log_probs, torch.Tensor)
+        and log_probs.dim() >= 1
+        and log_probs.shape[0] == batch_size
+    ):
+        log_probs = log_probs[sample_idx].contiguous()
+    debug_tensors = None
+    if rtd.rollout_debug_tensors:
+        rd = rtd.rollout_debug_tensors
+        debug_tensors = RolloutDebugTensors(
+            rollout_variance_noises=_extract_single_sample_tensor(
+                rd.rollout_variance_noises, sample_idx, batch_size
+            ),
+            rollout_prev_sample_means=_extract_single_sample_tensor(
+                rd.rollout_prev_sample_means, sample_idx, batch_size
+            ),
+            rollout_noise_std_devs=_extract_single_sample_tensor(
+                rd.rollout_noise_std_devs, sample_idx, batch_size
+            ),
+            rollout_model_outputs=_extract_single_sample_tensor(
+                rd.rollout_model_outputs, sample_idx, batch_size
+            ),
+        )
+    denoising_env = None
+    if rtd.denoising_env:
+        env = rtd.denoising_env
+        denoising_env = RolloutDenoisingEnv(
+            image_kwargs=(
+                _extract_single_sample_tensor(env.image_kwargs, sample_idx, batch_size)
+                if env.image_kwargs
+                else None
+            ),
+            pos_cond_kwargs=(
+                _extract_single_sample_tensor(
+                    env.pos_cond_kwargs, sample_idx, batch_size
+                )
+                if env.pos_cond_kwargs
+                else None
+            ),
+            neg_cond_kwargs=(
+                _extract_single_sample_tensor(
+                    env.neg_cond_kwargs, sample_idx, batch_size
+                )
+                if env.neg_cond_kwargs
+                else None
+            ),
+            guidance=(
+                _extract_single_sample_tensor(env.guidance, sample_idx, batch_size)
+                if env.guidance is not None
+                else None
+            ),
+        )
+    dit_trajectory = None
+    if rtd.dit_trajectory:
+        dit = rtd.dit_trajectory
+        dit_trajectory = RolloutDitTrajectory(
+            latents=_extract_single_sample_tensor(dit.latents, sample_idx, batch_size),
+            timesteps=dit.timesteps,
+        )
+    return RolloutTrajectoryData(
+        rollout_log_probs=log_probs,
+        rollout_debug_tensors=debug_tensors,
+        denoising_env=denoising_env,
+        dit_trajectory=dit_trajectory,
+    )
+
+
+def _serialize_rollout_trajectory(
+    rtd: RolloutTrajectoryData | None,
+    *,
+    serialized_dit_timesteps: 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:
+        return None, None, None, None
+    serialized_log_probs = _maybe_serialize(rtd.rollout_log_probs)
+    serialized_debug_tensors = None
+    if rtd.rollout_debug_tensors:
+        rd = rtd.rollout_debug_tensors
+        serialized_debug_tensors = {
+            "rollout_variance_noises": _maybe_serialize(rd.rollout_variance_noises),
+            "rollout_prev_sample_means": _maybe_serialize(rd.rollout_prev_sample_means),
+            "rollout_noise_std_devs": _maybe_serialize(rd.rollout_noise_std_devs),
+            "rollout_model_outputs": _maybe_serialize(rd.rollout_model_outputs),
+        }
+    serialized_denoising_env = None
+    if rtd.denoising_env:
+        env = rtd.denoising_env
+        serialized_denoising_env = {
+            "image_kwargs": (
+                _maybe_serialize(env.image_kwargs) if env.image_kwargs else None
+            ),
+            "pos_cond_kwargs": (
+                _maybe_serialize(env.pos_cond_kwargs) if env.pos_cond_kwargs else None
+            ),
+            "neg_cond_kwargs": (
+                _maybe_serialize(env.neg_cond_kwargs) if env.neg_cond_kwargs else None
+            ),
+            "guidance": (
+                _maybe_serialize(env.guidance) if env.guidance is not None else None
+            ),
+        }
+    serialized_dit_trajectory = None
+    if rtd.dit_trajectory:
+        dit = rtd.dit_trajectory
+        serialized_dit_trajectory = {
+            "latents": (
+                _maybe_serialize(dit.latents) if dit.latents is not None else None
+            ),
+            "timesteps": serialized_dit_timesteps,
+        }
+    return (
+        serialized_log_probs,
+        serialized_debug_tensors,
+        serialized_denoising_env,
+        serialized_dit_trajectory,
+    )
+
+
+def _build_response(
+    request_id: str, prompt: str, seed: int, rollout: bool, result: OutputBatch
+) -> list[RolloutResponse]:
+    """
+    rollout: bool - set to False when evaluating the model
+    """
+    batch_size = result.output.shape[0]
+    inference_time_s = (
+        result.metrics.total_duration_s
+        if result.metrics and result.metrics.total_duration_s > 0
+        else None
+    )
+    peak_memory_mb = result.peak_memory_mb if result.peak_memory_mb > 0 else None
+    rollout_trajectory_data = result.rollout_trajectory_data
+    if rollout:
+        assert (
+            rollout_trajectory_data is not None
+        ), "rollout_trajectory_data must be present when rollout=True"
+
+    serialized_dit_timesteps = None
+    if rollout and rollout_trajectory_data and rollout_trajectory_data.dit_trajectory:
+        serialized_dit_timesteps = _maybe_serialize(
+            rollout_trajectory_data.dit_trajectory.timesteps
+        )
+
+    responses: list[RolloutResponse] = []
+    for sample_idx in range(batch_size):
+        out_i = result.output[sample_idx].contiguous()
+        serialized_generated_output = _maybe_serialize(out_i)
+        if not rollout:
+            responses.append(
+                RolloutResponse(
+                    request_id=request_id,
+                    prompt=prompt,
+                    seed=seed,
+                    generated_output=serialized_generated_output,
+                    inference_time_s=inference_time_s,
+                    peak_memory_mb=peak_memory_mb,
+                )
+            )
+            continue
+        per_sample_trajectory = _slice_rollout_trajectory_for_sample(
+            result.rollout_trajectory_data, sample_idx, batch_size
+        )
+        (
+            serialized_log_probs,
+            serialized_debug_tensors,
+            serialized_denoising_env,
+            serialized_dit_trajectory,
+        ) = _serialize_rollout_trajectory(
+            per_sample_trajectory,
+            serialized_dit_timesteps=serialized_dit_timesteps,
+        )
+        responses.append(
+            RolloutResponse(
+                request_id=request_id,
+                prompt=prompt,
+                seed=seed,
+                generated_output=serialized_generated_output,
+                rollout_log_probs=serialized_log_probs,
+                rollout_debug_tensors=serialized_debug_tensors,
+                denoising_env=serialized_denoising_env,
+                dit_trajectory=serialized_dit_trajectory,
+                inference_time_s=inference_time_s,
+                peak_memory_mb=peak_memory_mb,
+            )
+        )
+    return responses
+
+
+@router.post("/generate", response_model=list[RolloutResponse])
+async def rollout_generate(request: RolloutRequest):
+    request_id = generate_request_id()
+    server_args = get_global_server_args()
+    sampling_kwargs: dict = dict(
+        prompt=request.prompt,
+        negative_prompt=request.negative_prompt,
+        seed=request.seed,
+        generator_device=request.generator_device,
+        width=request.width,
+        height=request.height,
+        num_inference_steps=request.num_inference_steps,
+        num_outputs_per_prompt=request.num_outputs_per_prompt,
+        guidance_scale=request.guidance_scale,
+        true_cfg_scale=request.true_cfg_scale,
+        num_frames=request.num_frames,
+        fps=request.fps,
+        image_path=request.image_path,
+        rollout=request.rollout,
+        rollout_sde_type=request.rollout_sde_type,
+        rollout_noise_level=request.rollout_noise_level,
+        rollout_log_prob_no_const=request.rollout_log_prob_no_const,
+        rollout_debug_mode=request.rollout_debug_mode,
+        rollout_return_denoising_env=request.rollout_return_denoising_env,
+        rollout_return_dit_trajectory=request.rollout_return_dit_trajectory,
+        suppress_logs=request.suppress_logs,
+        save_output=False,
+        return_trajectory_latents=False,
+        return_trajectory_decoded=False,
+    )
+    if request.extra_sampling_params:
+        sampling_kwargs.update(request.extra_sampling_params)
+        sampling_kwargs["rollout"] = request.rollout
+    sampling_kwargs = {k: v for k, v in sampling_kwargs.items() if v is not None}
+    try:
+        sampling_params = build_sampling_params(request_id, **sampling_kwargs)
+    except Exception as exc:
+        raise HTTPException(
+            status_code=400, detail=f"Invalid sampling params: {exc}"
+        ) from exc
+    pipeline_request = prepare_request(
+        server_args=server_args, sampling_params=sampling_params
+    )
+    try:
+        output_batch: OutputBatch = await async_scheduler_client.forward(
+            pipeline_request
+        )
+    except Exception as exc:
+        logger.error("Rollout generation failed: %s", exc, exc_info=True)
+        raise HTTPException(
+            status_code=500, detail=f"Generation failed: {exc}"
+        ) from exc
+    if output_batch.error:
+        raise HTTPException(status_code=500, detail=output_batch.error)
+    rollout_responses = _build_response(
+        request_id, request.prompt, request.seed, request.rollout, output_batch
+    )
+    return ORJSONResponse(content=[r.model_dump() for r in rollout_responses])

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

@@ -0,0 +1,45 @@
+"""Tensor serialization for post-training / rollout HTTP responses."""
+
+from __future__ import annotations
+
+import base64
+from typing import Any
+
+import torch
+from safetensors.torch import load, save
+
+
+def tensor_to_base64(t: torch.Tensor) -> str:
+    t = t.detach().contiguous().cpu()
+    raw = save({"t": t})
+    return base64.b64encode(raw).decode("ascii")
+
+
+def base64_to_tensor(s: str) -> torch.Tensor:
+    raw = base64.b64decode(s)
+    return load(raw)["t"]
+
+
+def _maybe_serialize(obj: Any) -> Any:
+    if isinstance(obj, torch.Tensor):
+        return {
+            "__tensor__": True,
+            "data": tensor_to_base64(obj),
+            "shape": list(obj.shape),
+            "dtype": str(obj.dtype),
+        }
+    if isinstance(obj, dict):
+        return {k: _maybe_serialize(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_maybe_serialize(v) for v in obj]
+    return obj
+
+
+def _maybe_deserialize(obj: Any) -> Any:
+    if isinstance(obj, dict):
+        if obj.get("__tensor__"):
+            return base64_to_tensor(obj["data"])
+        return {k: _maybe_deserialize(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_maybe_deserialize(v) for v in obj]
+    return obj

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

@@ -82,11 +82,8 @@ from sglang.multimodal_gen.runtime.platforms import (
     AttentionBackendEnum,
     current_platform,
 )
-from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import (
-    RolloutTrajectoryData,
-)
-from sglang.multimodal_gen.runtime.post_training.scheduler_rl_mixin import (
-    SchedulerRLMixin,
+from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import (
+    RolloutDenoisingMixin,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
@@ -147,7 +144,7 @@ class DenoisingStepState:
     attn_metadata: Any | None
 
 
-class DenoisingStage(PipelineStage):
+class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
     """
     Stage for running the denoising loop in diffusion pipelines.
 
@@ -192,43 +189,6 @@ class DenoisingStage(PipelineStage):
         self._cached_num_steps = None
         self._is_warmed_up = False
 
-    def _maybe_prepare_rollout(self, batch: Req):
-        """Prepare denoising loop for rollout."""
-        if not isinstance(self.scheduler, SchedulerRLMixin):
-            if batch.rollout:
-                raise ValueError(
-                    f"Scheduler {type(self.scheduler)} does not support rollout"
-                )
-            return
-
-        self.scheduler.release_rollout_resources(batch)
-        if batch.rollout:
-            self.scheduler.prepare_rollout(
-                batch=batch,
-                pipeline_config=self.server_args.pipeline_config,
-            )
-
-    def _maybe_collect_rollout_log_probs(self, batch: Req):
-        """Get rollout log probs and store into batch for reward calculation."""
-        if not isinstance(self.scheduler, SchedulerRLMixin):
-            if batch.rollout:
-                raise ValueError(
-                    f"Scheduler {type(self.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)
-            )
-            if getattr(batch, "rollout_debug_mode", False):
-                batch.rollout_trajectory_data.rollout_debug_tensors = (
-                    self.scheduler.collect_rollout_debug_tensors(batch)
-                )
-            self.scheduler.release_rollout_resources(batch)
-
     def _maybe_enable_torch_compile(self, module: object) -> None:
         """
         Compile a module with torch.compile, and enable inductor overlap tweak if available.
@@ -901,10 +861,6 @@ class DenoisingStage(PipelineStage):
             trajectory_tensor = None
             trajectory_timesteps_tensor = None
 
-        # Gather log probs for rollout
-        if batch.rollout:
-            self._maybe_collect_rollout_log_probs(batch)
-
         # Gather results if using sequence parallelism
         latents, trajectory_tensor = self._postprocess_sp_latents(
             batch, latents, trajectory_tensor
@@ -1137,6 +1093,15 @@ class DenoisingStage(PipelineStage):
         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,
+            )
         denoising_start_time = time.time()
         self._before_denoising_loop(ctx, batch, server_args)
         # to avoid device-sync caused by timestep comparison
@@ -1164,6 +1129,17 @@ class DenoisingStage(PipelineStage):
                             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 batch.rollout:
+                            self._maybe_append_dit_trajectory_step(
+                                batch=batch,
+                                latents=ctx.latents,
+                                timestep_value=step.t_host,
+                            )
                         self._run_denoising_step(ctx, step, batch, server_args)
                         self._record_trajectory(ctx, step, batch, server_args)
 
@@ -1185,6 +1161,16 @@ class DenoisingStage(PipelineStage):
                 (denoising_end_time - denoising_start_time) / len(ctx.timesteps),
             )
 
+        # Rollout postprocessing must run BEFORE _finalize_denoising_loop so
+        # the final scheduler.step output (ctx.latents) is still SP-sharded and
+        # can be gathered uniformly alongside the per-step dit_trajectory via
+        # gather_stacked_latents_for_sp.
+        if batch.rollout:
+            self._postprocess_rollout_outputs(
+                batch=batch,
+                latents=ctx.latents,
+                server_args=server_args,
+            )
         self._finalize_denoising_loop(ctx, batch, server_args)
         return batch
 

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

@@ -39,8 +39,24 @@ class RolloutDebugTensors:
 
 
 @dataclass
-class RolloutTrajectoryData:
-    """Container for rollout-specific trajectory outputs."""
+class RolloutDenoisingEnv:
+    image_kwargs: dict[str, Any] | None = None
+    pos_cond_kwargs: dict[str, Any] | None = None
+    neg_cond_kwargs: dict[str, Any] | None = None
+    guidance: torch.Tensor | None = None
 
+
+@dataclass
+class RolloutDitTrajectory:
+    # [B, T+1, ...]: per-step noisy latents x_{t_0..t_{T-1}} followed by the
+    # final denoised latent x_{t_T} (last scheduler.step output).
+    latents: torch.Tensor | None = None
+    timesteps: torch.Tensor | None = None  # [T]
+
+
+@dataclass
+class RolloutTrajectoryData:
     rollout_log_probs: torch.Tensor | None = None
     rollout_debug_tensors: RolloutDebugTensors | None = None
+    denoising_env: RolloutDenoisingEnv | None = None
+    dit_trajectory: RolloutDitTrajectory | None = None

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

@@ -0,0 +1,197 @@
+"""Mixin for rollout-related denoising hooks.
+
+Moved out of DenoisingStage to keep the core stage lean.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import (
+    RolloutDenoisingEnv,
+    RolloutDitTrajectory,
+    RolloutTrajectoryData,
+)
+from sglang.multimodal_gen.runtime.post_training.scheduler_rl_mixin import (
+    SchedulerRLMixin,
+)
+from sglang.multimodal_gen.runtime.post_training.sp_utils import (
+    gather_stacked_latents_for_sp,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+def _kwargs_to_cpu(d: Any) -> Any:
+    if isinstance(d, torch.Tensor):
+        return d.detach().cpu()
+    if isinstance(d, dict):
+        return {k: _kwargs_to_cpu(v) for k, v in d.items()}
+    if isinstance(d, list):
+        return [_kwargs_to_cpu(v) for v in d]
+    if isinstance(d, tuple):
+        return tuple(_kwargs_to_cpu(v) for v in d)
+    return d
+
+
+class RolloutDenoisingMixin:
+
+    def _maybe_prepare_rollout(self, batch: Req):
+        """Prepare denoising loop for rollout."""
+        if not isinstance(self.scheduler, SchedulerRLMixin):
+            if batch.rollout:
+                raise ValueError(
+                    f"Scheduler {type(self.scheduler)} does not support rollout"
+                )
+            return
+
+        self.scheduler.release_rollout_resources(batch)
+        if batch.rollout:
+            self.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):
+            if batch.rollout:
+                raise ValueError(
+                    f"Scheduler {type(self.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)
+            )
+            if batch.rollout_debug_mode:
+                batch.rollout_trajectory_data.rollout_debug_tensors = (
+                    self.scheduler.collect_rollout_debug_tensors(batch)
+                )
+            self.scheduler.release_rollout_resources(batch)
+
+    def _postprocess_rollout_outputs(
+        self,
+        batch: Req,
+        latents: torch.Tensor,
+        server_args: ServerArgs,
+    ) -> None:
+        """Finalize rollout-only outputs.
+
+        Must be called before ``_post_denoising_loop`` so that ``latents`` (the
+        last ``scheduler.step`` output) is still SP-sharded and can be gathered
+        uniformly with the per-step trajectory latents.
+        """
+        self._maybe_collect_rollout_log_probs(batch)
+        # Append the final denoised latent as the (T+1)-th entry of the
+        # dit-trajectory latents list.
+        state = getattr(batch, "_rollout_dit_env_state", None)
+        if state is not None and batch.rollout and batch.rollout_return_dit_trajectory:
+            state["step_latents"].append(latents.detach())
+        self._maybe_finalize_dit_env_collection(
+            batch=batch,
+            pipeline_config=server_args.pipeline_config,
+        )
+
+    def _maybe_init_denoising_env_collection(
+        self,
+        batch,
+        pipeline_config,
+        image_kwargs: dict[str, Any],
+        pos_cond_kwargs: dict[str, Any],
+        neg_cond_kwargs: dict[str, Any],
+        guidance: torch.Tensor | None,
+    ) -> None:
+        collect_env = batch.rollout_return_denoising_env
+        collect_traj = batch.rollout_return_dit_trajectory
+        if not (collect_env or collect_traj):
+            batch._rollout_dit_env_state = None
+            return
+
+        sanitize = getattr(pipeline_config, "sanitize_dit_env_kwargs", lambda x: x)
+        if collect_env:
+            env = RolloutDenoisingEnv(
+                image_kwargs=_kwargs_to_cpu(sanitize(image_kwargs)),
+                pos_cond_kwargs=_kwargs_to_cpu(sanitize(pos_cond_kwargs)),
+                neg_cond_kwargs=(
+                    _kwargs_to_cpu(sanitize(neg_cond_kwargs))
+                    if neg_cond_kwargs
+                    else None
+                ),
+                guidance=guidance.detach().cpu() if guidance is not None else None,
+            )
+            pos_src = pos_cond_kwargs
+            neg_src = neg_cond_kwargs
+        else:
+            env = None
+            pos_src = None
+            neg_src = None
+
+        batch._rollout_dit_env_state = {
+            "env": env,
+            "step_latents": [],
+            "step_timesteps": [],
+            "pos_cond_kwargs_src": pos_src,
+            "neg_cond_kwargs_src": neg_src,
+        }
+
+    def _maybe_append_dit_trajectory_step(
+        self,
+        batch,
+        latents: torch.Tensor,
+        timestep_value: torch.Tensor,
+    ) -> None:
+        if not batch.rollout or not batch.rollout_return_dit_trajectory:
+            return
+        state = getattr(batch, "_rollout_dit_env_state", None)
+        if state is None:
+            return
+
+        state["step_latents"].append(latents.detach())
+        state["step_timesteps"].append(timestep_value.detach().cpu())
+
+    def _maybe_finalize_dit_env_collection(self, batch, pipeline_config) -> None:
+        state = getattr(batch, "_rollout_dit_env_state", None)
+        if state is None:
+            return
+
+        env: RolloutDenoisingEnv | None = state["env"]
+        step_latents: list[torch.Tensor] = state["step_latents"]
+        step_timesteps: list[torch.Tensor] = state["step_timesteps"]
+
+        if batch.rollout_trajectory_data is None:
+            batch.rollout_trajectory_data = RolloutTrajectoryData()
+
+        if step_latents and batch.rollout_return_dit_trajectory:
+            step_latents_tensor = torch.stack(step_latents, dim=1)
+            step_latents_tensor = gather_stacked_latents_for_sp(
+                pipeline_config=pipeline_config,
+                batch=batch,
+                stacked_latents=step_latents_tensor,
+            )
+            batch.rollout_trajectory_data.dit_trajectory = RolloutDitTrajectory(
+                latents=step_latents_tensor.cpu(),
+                timesteps=torch.stack(step_timesteps, dim=0).cpu(),
+            )
+
+        if env is not None and batch.rollout_return_denoising_env:
+            sanitize = getattr(pipeline_config, "sanitize_dit_env_kwargs", lambda x: x)
+            gather_fn = getattr(pipeline_config, "gather_dit_env_static_for_sp", None)
+
+            pos_src = state.get("pos_cond_kwargs_src")
+            if pos_src is not None and env.pos_cond_kwargs is not None:
+                gathered_pos = gather_fn(batch, pos_src) if gather_fn else pos_src
+                env.pos_cond_kwargs = _kwargs_to_cpu(sanitize(gathered_pos))
+
+            neg_src = state.get("neg_cond_kwargs_src")
+            if neg_src is not None and env.neg_cond_kwargs is not None:
+                gathered_neg = gather_fn(batch, neg_src) if gather_fn else neg_src
+                env.neg_cond_kwargs = _kwargs_to_cpu(sanitize(gathered_neg))
+
+            batch.rollout_trajectory_data.denoising_env = env
+
+        batch._rollout_dit_env_state = None

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

@@ -33,7 +33,8 @@ class SchedulerRLDebugMixin:
     ) -> None:
         rollout_session_data = batch._rollout_session_data
         batch_size = variance_noise.shape[0]
-        rollout_session_data.local_variance_noises.append(variance_noise)
+        # the underlying noise buffer in batch._rollout_session_data.noise_buffer is reused.
+        rollout_session_data.local_variance_noises.append(variance_noise.clone())
         rollout_session_data.local_prev_sample_means.append(prev_sample_mean)
         rollout_session_data.local_noise_std_devs.append(
             noise_std_dev.expand((batch_size, 1))
@@ -86,13 +87,13 @@ class SchedulerRLDebugMixin:
 
             # Gather on packed tensors first.
             variance_noises_packed = pipeline_config.gather_latents_for_sp(
-                variance_noises_packed
+                variance_noises_packed, batch=batch
             )
             prev_sample_means_packed = pipeline_config.gather_latents_for_sp(
-                prev_sample_means_packed
+                prev_sample_means_packed, batch=batch
             )
             model_outputs_packed = pipeline_config.gather_latents_for_sp(
-                model_outputs_packed
+                model_outputs_packed, batch=batch
             )
 
             # Unpack back to [B, T, ...].

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

@@ -7,12 +7,8 @@ from typing import Any, Union
 import torch
 
 from sglang.multimodal_gen.runtime.distributed import (
-    get_local_torch_device,
     get_sp_world_size,
 )
-from sglang.multimodal_gen.runtime.distributed.communication_op import (
-    sequence_model_parallel_all_reduce,
-)
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import (
     RolloutSessionData,
@@ -107,7 +103,7 @@ class SchedulerRLMixin(SchedulerRLDebugMixin):
             )
 
         sharded_noise, _ = rollout_session_data.pipeline_config.shard_latents_for_sp(
-            batch, buffer
+            batch=batch, latents=buffer
         )
         if tuple(sharded_noise.shape) != local_shape:
             raise ValueError(
@@ -145,10 +141,19 @@ class SchedulerRLMixin(SchedulerRLDebugMixin):
             ), "True log-probability computation requires a non-zero noise level."
 
         dt = next_sigma - current_sigma
+        # sde/cps: cast to fp32 to match flowGRPO semantics and avoid the
+        # 0-dim-fp32 wrapped-scalar promotion demoting log-prob to bf16.
+        # ode: keep dtypes unchanged so rollout(ode) stays bit-exact with
+        # rollout=False (scheduling_flow_match_euler_discrete.step()).
+        # log_prob is computed on the full pre-shard noise buffer so SP ranks
+        # produce identical sums — see collect_rollout_log_probs().
         if sde_type == "sde":
+            model_output = model_output.float()
+            sample = sample.float()
             variance_noise = self._rollout_variance_noise(
                 batch, model_output, generator
             )
+            full_variance_noise = rollout_session_data.noise_buffer
             std_dev_t = (
                 torch.sqrt(
                     current_sigma
@@ -173,12 +178,15 @@ class SchedulerRLMixin(SchedulerRLDebugMixin):
 
             weighted_variance_noise = variance_noise * noise_std_dev
             prev_sample = prev_sample_mean + weighted_variance_noise
-            log_prob_no_const_val = -(weighted_variance_noise**2)
+            log_prob_no_const_val = -((full_variance_noise * noise_std_dev) ** 2)
 
         elif sde_type == "cps":
+            model_output = model_output.float()
+            sample = sample.float()
             variance_noise = self._rollout_variance_noise(
                 batch, model_output, generator
             )
+            full_variance_noise = rollout_session_data.noise_buffer
             std_dev_t = next_sigma * math.sin(noise_level * math.pi / 2)
             noise_std_dev = std_dev_t
             pred_original_sample = sample - current_sigma * model_output
@@ -189,7 +197,7 @@ class SchedulerRLMixin(SchedulerRLDebugMixin):
 
             weighted_variance_noise = variance_noise * noise_std_dev
             prev_sample = prev_sample_mean + weighted_variance_noise
-            log_prob_no_const_val = -(weighted_variance_noise**2)
+            log_prob_no_const_val = -((full_variance_noise * noise_std_dev) ** 2)
 
         elif sde_type == "ode":
             prev_sample = sample + dt * model_output
@@ -198,7 +206,11 @@ class SchedulerRLMixin(SchedulerRLDebugMixin):
             noise_std_dev = torch.zeros(
                 (), device=model_output.device, dtype=model_output.dtype
             )
-            log_prob_no_const_val = torch.zeros_like(model_output)
+            log_prob_no_const_val = torch.zeros(
+                rollout_session_data.latents_shape,
+                device=model_output.device,
+                dtype=torch.float32,
+            )
             assert (
                 log_prob_no_const
             ), "p_ode is always 0, true log_prob is meaningless, set rollout_log_prob_no_const to True to enable log_prob computation"
@@ -245,25 +257,20 @@ class SchedulerRLMixin(SchedulerRLDebugMixin):
         self, batch
     ) -> tuple[torch.Tensor, torch.Tensor]:
         rollout_session_data = self._get_rollout_session_data(batch)
-        values_sum = torch.stack(rollout_session_data.local_log_prob_sum, dim=-1)
-        values_count = torch.stack(rollout_session_data.local_log_prob_count, dim=-1)
+        # [B, T]: batch dim 0, denoising step dim 1
+        values_sum = torch.stack(rollout_session_data.local_log_prob_sum, dim=1)
+        values_count = torch.stack(rollout_session_data.local_log_prob_count, dim=1)
         rollout_session_data.local_log_prob_sum = []
         rollout_session_data.local_log_prob_count = []
         return values_sum, values_count
 
     def collect_rollout_log_probs(self, batch: Req) -> torch.Tensor | None:
-        """Consume local rollout log probs and merge for all SP ranks."""
+        """Per-step sums are already computed on the full pre-shard noise
+        buffer inside flow_sde_sampling, so every SP rank holds identical
+        values here and no all-reduce is needed."""
 
         trajectory_log_prob_sum, trajectory_log_prob_count = (
             self.consume_local_rollout_log_probs(batch)
         )
-        if get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False):
-            packed = torch.stack(
-                [trajectory_log_prob_sum, trajectory_log_prob_count], dim=0
-            ).to(get_local_torch_device())
-            sequence_model_parallel_all_reduce(packed)
-            trajectory_log_prob_sum = packed[0]
-            trajectory_log_prob_count = packed[1]
-
         rollout_log_probs_tensor = trajectory_log_prob_sum / trajectory_log_prob_count
         return rollout_log_probs_tensor.cpu()

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

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Sequence Parallel helpers for post-training rollout code."""
+
+from __future__ import annotations
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed import (
+    get_local_torch_device,
+    get_sp_world_size,
+)
+from sglang.multimodal_gen.runtime.distributed.communication_op import (
+    sequence_model_parallel_all_gather,
+    sequence_model_parallel_all_reduce,
+)
+
+
+def should_do_sp_collective(batch) -> bool:
+    return get_sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False)
+
+
+def gather_stacked_latents_for_sp(
+    pipeline_config,
+    batch,
+    stacked_latents: torch.Tensor,
+) -> torch.Tensor:
+    if not should_do_sp_collective(batch):
+        return stacked_latents
+    if stacked_latents.dim() < 2:
+        return stacked_latents
+    bsz, t_steps = stacked_latents.shape[0], stacked_latents.shape[1]
+    flat_inputs = stacked_latents.flatten(0, 1).contiguous()
+    gathered_flat_inputs = pipeline_config.gather_latents_for_sp(
+        flat_inputs, batch=batch
+    )
+    return gathered_flat_inputs.unflatten(0, (bsz, t_steps))
+
+
+def all_reduce_if_sp_sharded(batch, tensor: torch.Tensor) -> torch.Tensor:
+    if not should_do_sp_collective(batch):
+        return tensor
+    tensor = tensor.to(get_local_torch_device())
+    sequence_model_parallel_all_reduce(tensor)
+    return tensor
+
+
+def all_gather_if_sp_sharded(batch, x: torch.Tensor, dim: int = 0) -> torch.Tensor:
+    if not should_do_sp_collective(batch):
+        return x
+    x = x.to(get_local_torch_device()).contiguous()
+    return sequence_model_parallel_all_gather(x, dim=dim)
+
+
+def maybe_trim_sp_rope_seq_for_batch(batch, rope: torch.Tensor) -> torch.Tensor:
+    raw = getattr(batch, "raw_latent_shape", None)
+    if raw is None or len(raw) < 2:
+        return rope
+    target = int(raw[1])
+    if rope.shape[0] > target:
+        return rope[:target]
+    return rope

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

@@ -0,0 +1,360 @@
+"""Unit tests for the rollout generate API (serialization, io_struct, rollout_api)."""
+
+import types
+import unittest
+
+import torch
+
+from sglang.multimodal_gen.runtime.entrypoints.post_training.utils import (
+    _maybe_deserialize,
+    _maybe_serialize,
+    base64_to_tensor,
+    tensor_to_base64,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
+from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import (
+    RolloutDebugTensors,
+    RolloutDenoisingEnv,
+    RolloutDitTrajectory,
+    RolloutTrajectoryData,
+)
+
+
+class TestTensorToBase64Roundtrip(unittest.TestCase):
+
+    def _roundtrip(self, t: torch.Tensor):
+        encoded = tensor_to_base64(t)
+        self.assertIsInstance(encoded, str)
+        decoded = base64_to_tensor(encoded)
+        self.assertTrue(
+            torch.equal(t, decoded), f"Mismatch for shape={t.shape} dtype={t.dtype}"
+        )
+
+    def test_float32_1d(self):
+        self._roundtrip(torch.randn(16))
+
+    def test_float32_nd(self):
+        self._roundtrip(torch.randn(2, 4, 8, 8))
+
+    def test_float16(self):
+        self._roundtrip(torch.randn(3, 5).half())
+
+    def test_int64(self):
+        self._roundtrip(torch.arange(10))
+
+    def test_bool(self):
+        self._roundtrip(torch.tensor([True, False, True]))
+
+    def test_scalar(self):
+        self._roundtrip(torch.tensor(3.14))
+
+    def test_empty(self):
+        self._roundtrip(torch.empty(0))
+
+    def test_cuda_tensor_moves_to_cpu(self):
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA not available")
+        t = torch.randn(4, device="cuda")
+        encoded = tensor_to_base64(t)
+        decoded = base64_to_tensor(encoded)
+        self.assertTrue(torch.equal(t.cpu(), decoded))
+
+    def test_non_contiguous(self):
+        t = torch.randn(4, 6)[:, ::2]
+        self.assertFalse(t.is_contiguous())
+        self._roundtrip(t.contiguous())
+        decoded = base64_to_tensor(tensor_to_base64(t))
+        self.assertTrue(torch.equal(t.contiguous(), decoded))
+
+    def test_grad_tensor_detaches(self):
+        t = torch.randn(3, requires_grad=True)
+        encoded = tensor_to_base64(t)
+        decoded = base64_to_tensor(encoded)
+        self.assertFalse(decoded.requires_grad)
+        self.assertTrue(torch.equal(t.detach(), decoded))
+
+
+class TestMaybeSerialize(unittest.TestCase):
+    def test_tensor(self):
+        t = torch.randn(2, 3)
+        result = _maybe_serialize(t)
+        self.assertIsInstance(result, dict)
+        self.assertTrue(result["__tensor__"])
+        self.assertEqual(result["shape"], [2, 3])
+        self.assertEqual(result["dtype"], "torch.float32")
+        decoded = base64_to_tensor(result["data"])
+        self.assertTrue(torch.equal(t, decoded))
+
+    def test_dict_with_tensors(self):
+        d = {"a": torch.tensor([1.0]), "b": "hello", "c": 42}
+        result = _maybe_serialize(d)
+        self.assertIsInstance(result, dict)
+        self.assertTrue(result["a"]["__tensor__"])
+        self.assertEqual(result["b"], "hello")
+        self.assertEqual(result["c"], 42)
+
+    def test_list_with_tensors(self):
+        lst = [torch.tensor(1.0), "text", torch.tensor(2.0)]
+        result = _maybe_serialize(lst)
+        self.assertIsInstance(result, list)
+        self.assertTrue(result[0]["__tensor__"])
+        self.assertEqual(result[1], "text")
+        self.assertTrue(result[2]["__tensor__"])
+
+    def test_nested_structure(self):
+        nested = {
+            "level1": {"level2": [torch.tensor(1.0), {"level3": torch.tensor(2.0)}]}
+        }
+        result = _maybe_serialize(nested)
+        self.assertTrue(result["level1"]["level2"][0]["__tensor__"])
+        self.assertTrue(result["level1"]["level2"][1]["level3"]["__tensor__"])
+
+    def test_none_passthrough(self):
+        self.assertIsNone(_maybe_serialize(None))
+
+    def test_plain_values_passthrough(self):
+        self.assertEqual(_maybe_serialize(42), 42)
+        self.assertEqual(_maybe_serialize("hello"), "hello")
+        self.assertAlmostEqual(_maybe_serialize(3.14), 3.14)
+
+    def test_tuple_becomes_list(self):
+        result = _maybe_serialize((torch.tensor(1.0), 2))
+        self.assertIsInstance(result, list)
+        self.assertEqual(len(result), 2)
+
+
+from sglang.multimodal_gen.runtime.entrypoints.post_training.rollout_api import (
+    _build_response,
+    _serialize_rollout_trajectory,
+)
+
+
+class TestSerializeRolloutTrajectory(unittest.TestCase):
+    def test_none_input(self):
+        log_probs, debug, env, dit_traj = _serialize_rollout_trajectory(None)
+        self.assertIsNone(log_probs)
+        self.assertIsNone(debug)
+        self.assertIsNone(env)
+        self.assertIsNone(dit_traj)
+
+    def test_log_probs_only(self):
+        rtd = RolloutTrajectoryData(
+            rollout_log_probs=torch.tensor([-1.0, -2.0]),
+        )
+        log_probs, debug, env, dit_traj = _serialize_rollout_trajectory(rtd)
+        self.assertIsNotNone(log_probs)
+        self.assertTrue(log_probs["__tensor__"])
+        self.assertIsNone(debug)
+        self.assertIsNone(env)
+        self.assertIsNone(dit_traj)
+
+    def test_log_probs_none_in_rtd(self):
+        rtd = RolloutTrajectoryData(rollout_log_probs=None)
+        log_probs, debug, env, dit_traj = _serialize_rollout_trajectory(rtd)
+        self.assertIsNone(log_probs)
+        self.assertIsNone(debug)
+        self.assertIsNone(env)
+        self.assertIsNone(dit_traj)
+
+    def test_with_debug_tensors(self):
+        dt = RolloutDebugTensors(
+            rollout_variance_noises=torch.randn(2, 5, 4, 8, 8),
+            rollout_prev_sample_means=torch.randn(2, 5, 4, 8, 8),
+            rollout_noise_std_devs=torch.randn(2, 5, 1),
+            rollout_model_outputs=torch.randn(2, 5, 4, 8, 8),
+        )
+        rtd = RolloutTrajectoryData(
+            rollout_log_probs=torch.tensor([-0.5, -0.6]),
+            rollout_debug_tensors=dt,
+        )
+        log_probs, debug, env, dit_traj = _serialize_rollout_trajectory(rtd)
+        self.assertIsNotNone(log_probs)
+        self.assertIsNotNone(debug)
+        self.assertIsNone(env)
+        self.assertIsNone(dit_traj)
+        self.assertIn("rollout_variance_noises", debug)
+        self.assertIn("rollout_prev_sample_means", debug)
+        self.assertIn("rollout_noise_std_devs", debug)
+        self.assertIn("rollout_model_outputs", debug)
+        self.assertTrue(debug["rollout_variance_noises"]["__tensor__"])
+
+    def test_debug_tensors_with_none_fields(self):
+        dt = RolloutDebugTensors(
+            rollout_variance_noises=None,
+            rollout_prev_sample_means=torch.randn(1, 2, 4, 4, 4),
+            rollout_noise_std_devs=None,
+            rollout_model_outputs=None,
+        )
+        rtd = RolloutTrajectoryData(
+            rollout_log_probs=torch.tensor([-0.3]),
+            rollout_debug_tensors=dt,
+        )
+        log_probs, debug, env, dit_traj = _serialize_rollout_trajectory(rtd)
+        self.assertIsNotNone(debug)
+        self.assertIsNone(debug["rollout_variance_noises"])
+        self.assertTrue(debug["rollout_prev_sample_means"]["__tensor__"])
+        self.assertIsNone(env)
+        self.assertIsNone(dit_traj)
+
+    def test_with_denoising_env(self):
+        rtd = RolloutTrajectoryData(
+            denoising_env=RolloutDenoisingEnv(
+                image_kwargs={"encoder_hidden_states_image": [torch.randn(1, 8)]},
+                pos_cond_kwargs={"encoder_hidden_states": torch.randn(1, 8)},
+                neg_cond_kwargs={"encoder_hidden_states": torch.randn(1, 8)},
+                guidance=torch.tensor([3.5]),
+            ),
+            dit_trajectory=RolloutDitTrajectory(
+                latents=torch.randn(1, 5, 4, 2, 2, 2),
+                timesteps=torch.tensor([1.0, 0.75, 0.5, 0.25]),
+            ),
+        )
+        _, _, env, dit_traj = _serialize_rollout_trajectory(
+            rtd,
+            serialized_dit_timesteps=_maybe_serialize(rtd.dit_trajectory.timesteps),
+        )
+        self.assertIsNotNone(env)
+        self.assertIn("pos_cond_kwargs", env)
+        self.assertNotIn("trajectory", env)
+        self.assertIsNotNone(dit_traj)
+        self.assertIn("latents", dit_traj)
+        self.assertIn("timesteps", dit_traj)
+        self.assertTrue(dit_traj["latents"]["__tensor__"])
+        self.assertTrue(dit_traj["timesteps"]["__tensor__"])
+
+
+class TestBuildResponse(unittest.TestCase):
+    def _make_metrics(self, duration_s: float = 1.0):
+        return types.SimpleNamespace(total_duration_s=duration_s)
+
+    def test_minimal_output(self):
+        batch = OutputBatch(
+            output=torch.randn(1, 3, 1, 64, 64),
+            rollout_trajectory_data=RolloutTrajectoryData(
+                rollout_log_probs=torch.tensor([0.0]),
+            ),
+        )
+        batch.metrics = self._make_metrics(2.5)
+        resps = _build_response("r1", "prompt", 42, True, batch)
+        self.assertEqual(len(resps), 1)
+        resp = resps[0]
+        self.assertEqual(resp.request_id, "r1")
+        self.assertEqual(resp.prompt, "prompt")
+        self.assertEqual(resp.seed, 42)
+        self.assertIsNotNone(resp.generated_output)
+        self.assertIsNotNone(resp.rollout_log_probs)
+        lp = base64_to_tensor(resp.rollout_log_probs["data"])
+        self.assertEqual(lp.shape, ())
+        self.assertAlmostEqual(resp.inference_time_s, 2.5)
+
+    def test_full_response(self):
+        batch = OutputBatch(
+            output=torch.randn(1, 3, 1, 64, 64),
+            rollout_trajectory_data=RolloutTrajectoryData(
+                rollout_log_probs=torch.tensor([-0.5]),
+            ),
+            peak_memory_mb=8192.0,
+        )
+        batch.metrics = self._make_metrics(5.0)
+        resps = _build_response("r2", "test", 99, True, batch)
+        self.assertEqual(len(resps), 1)
+        resp = resps[0]
+        self.assertIsNotNone(resp.rollout_log_probs)
+        self.assertIsNone(resp.rollout_debug_tensors)
+        self.assertAlmostEqual(resp.peak_memory_mb, 8192.0)
+
+    def test_no_metrics(self):
+        batch = OutputBatch(
+            output=torch.randn(1, 3, 1, 64, 64),
+            rollout_trajectory_data=RolloutTrajectoryData(
+                rollout_log_probs=torch.tensor([0.0]),
+            ),
+        )
+        batch.metrics = None
+        resp = _build_response("r3", "p", 1, True, batch)[0]
+        self.assertIsNone(resp.inference_time_s)
+
+    def test_zero_metrics(self):
+        batch = OutputBatch(
+            output=torch.randn(1, 3, 1, 64, 64),
+            rollout_trajectory_data=RolloutTrajectoryData(
+                rollout_log_probs=torch.tensor([0.0]),
+            ),
+        )
+        batch.metrics = self._make_metrics(0.0)
+        resp = _build_response("r4", "p", 1, True, batch)[0]
+        self.assertIsNone(resp.inference_time_s)
+
+    def test_zero_peak_memory(self):
+        batch = OutputBatch(
+            output=torch.randn(1, 3, 1, 64, 64),
+            peak_memory_mb=0.0,
+            rollout_trajectory_data=RolloutTrajectoryData(
+                rollout_log_probs=torch.tensor([0.0]),
+            ),
+        )
+        batch.metrics = None
+        resp = _build_response("r6", "p", 1, True, batch)[0]
+        self.assertIsNone(resp.peak_memory_mb)
+
+    def test_batch_splits_log_probs_and_output(self):
+        B, T = 2, 3
+        batch = OutputBatch(
+            output=torch.randn(B, 1, 8, 8),
+            rollout_trajectory_data=RolloutTrajectoryData(
+                rollout_log_probs=torch.randn(B, T),
+            ),
+        )
+        batch.metrics = self._make_metrics(1.0)
+        resps = _build_response("rb", "p", 0, True, batch)
+        self.assertEqual(len(resps), B)
+        lp0 = base64_to_tensor(resps[0].rollout_log_probs["data"])
+        lp1 = base64_to_tensor(resps[1].rollout_log_probs["data"])
+        self.assertEqual(lp0.shape, (T,))
+        self.assertEqual(lp1.shape, (T,))
+        g0 = base64_to_tensor(resps[0].generated_output["data"])
+        g1 = base64_to_tensor(resps[1].generated_output["data"])
+        self.assertEqual(g0.shape, (1, 8, 8))
+        self.assertEqual(g1.shape, (1, 8, 8))
+        self.assertFalse(torch.equal(g0, g1))
+
+    def test_batch_dit_timesteps_on_each_row_one_serialize(self):
+        B, T, D = 2, 4, 3
+        batch = OutputBatch(
+            output=torch.randn(B, 1, 8, 8),
+            rollout_trajectory_data=RolloutTrajectoryData(
+                rollout_log_probs=torch.randn(B, T),
+                dit_trajectory=RolloutDitTrajectory(
+                    latents=torch.randn(B, T + 1, D),
+                    timesteps=torch.linspace(1.0, 0.0, T),
+                ),
+            ),
+        )
+        batch.metrics = self._make_metrics(1.0)
+        resps = _build_response("r", "p", 0, True, batch)
+        self.assertEqual(len(resps), B)
+        self.assertIsNotNone(resps[0].dit_trajectory)
+        self.assertIsNotNone(resps[1].dit_trajectory)
+        ts0 = base64_to_tensor(resps[0].dit_trajectory["timesteps"]["data"])
+        ts1 = base64_to_tensor(resps[1].dit_trajectory["timesteps"]["data"])
+        self.assertEqual(ts0.shape, (T,))
+        self.assertTrue(torch.equal(ts0, ts1))
+        self.assertEqual(
+            _maybe_deserialize(resps[1].dit_trajectory["latents"]).shape, (T + 1, D)
+        )
+
+    def test_rollout_false_omits_trajectory(self):
+        batch = OutputBatch(
+            output=torch.randn(2, 1, 8, 8),
+            rollout_trajectory_data=None,
+        )
+        batch.metrics = self._make_metrics(1.0)
+        resps = _build_response("r0", "p", 0, False, batch)
+        self.assertEqual(len(resps), 2)
+        self.assertIsNone(resps[0].rollout_log_probs)
+        self.assertIsNone(resps[1].rollout_log_probs)
+        self.assertIsNotNone(resps[0].generated_output)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -71,6 +71,48 @@ class TestSchedulerRolloutOdeUnit(unittest.TestCase):
         self.assertEqual(tuple(local_elem_count.shape), (sample.shape[0],))
         self.assertTrue(torch.all(local_elem_count == float(sample[0].numel())))
 
+    def test_ode_bit_exact_with_non_rollout_path(self):
+        """ODE rollout must produce the exact same prev_sample as the
+        non-rollout deterministic branch in
+        ``scheduling_flow_match_euler_discrete.step`` (``prev_sample =
+        sample + dt * model_output``). Uses bf16 model_output because the
+        wrapped-scalar promotion difference that a spurious
+        ``model_output.float()`` in the ODE branch would introduce is most
+        visible at bf16 precision."""
+        scheduler = _DummyScheduler()
+        batch = self._build_batch(debug_mode=False)
+        scheduler.prepare_rollout(batch)
+
+        sample = torch.randn(2, 4, 8, 8, dtype=torch.float32)
+        model_output = torch.randn_like(sample).to(torch.bfloat16)
+        current_sigma = torch.tensor(0.6, dtype=torch.float32)
+        next_sigma = torch.tensor(0.4, dtype=torch.float32)
+        dt = next_sigma - current_sigma
+
+        rollout_prev = scheduler.flow_sde_sampling(
+            batch,
+            model_output=model_output,
+            sample=sample,
+            current_sigma=current_sigma,
+            next_sigma=next_sigma,
+            generator=torch.Generator(device=sample.device).manual_seed(0),
+        )
+        # Exact expression used by the non-rollout branch at
+        # scheduling_flow_match_euler_discrete.py `prev_sample = sample +
+        # dt * model_output` (after the shared ``sample.to(fp32)`` cast).
+        non_rollout_prev = sample + dt * model_output
+
+        self.assertEqual(rollout_prev.dtype, non_rollout_prev.dtype)
+        self.assertTrue(torch.equal(rollout_prev, non_rollout_prev))
+        # Also verify the post-cast to model_output.dtype (what scheduler.step
+        # returns downstream) is bit-exact.
+        self.assertTrue(
+            torch.equal(
+                rollout_prev.to(model_output.dtype),
+                non_rollout_prev.to(model_output.dtype),
+            )
+        )
+
     def test_ode_debug_tensors_have_shape_safe_noise_std(self):
         scheduler = _DummyScheduler()
         batch = self._build_batch(debug_mode=True)
@@ -215,8 +257,19 @@ class TestSchedulerFlowGRPOStepAlignmentUnit(unittest.TestCase):
                 model_output = torch.randn(shape, generator=g, dtype=torch.float32)
                 sample = torch.randn(shape, generator=g, dtype=torch.float32)
                 variance_noise = torch.randn(shape, generator=g, dtype=torch.float32)
+
+                def _mock_rollout_variance_noise(_batch, *_args, **_kwargs):
+                    # flow_sde_sampling reads the full pre-shard noise from
+                    # rollout_session_data.noise_buffer to compute log_prob, so
+                    # the mock must populate it alongside returning the
+                    # (single-GPU trivially-sharded) noise.
+                    scheduler._get_rollout_session_data(  # type: ignore[attr-defined]
+                        _batch
+                    ).noise_buffer = variance_noise
+                    return variance_noise
+
                 scheduler._rollout_variance_noise = (  # type: ignore[method-assign]
-                    lambda _batch, *_args, **_kwargs: variance_noise
+                    _mock_rollout_variance_noise
                 )
 
                 prev_sgl = scheduler.flow_sde_sampling(
@@ -277,6 +330,57 @@ class TestSchedulerFlowGRPOStepAlignmentUnit(unittest.TestCase):
                         msg=f"{sde_type} seed={seed} {name} max_abs={err:.9f}",
                     )
 
+    def test_sde_cps_force_fp32_with_bf16_model_output(self):
+        """Regression for PyTorch's wrapped-scalar promotion trap: a 0-dim
+        fp32 ``noise_std_dev`` multiplied by an N-dim bf16 tensor silently
+        demotes to bf16, which would corrupt log-prob precision. SDE/CPS
+        branches therefore cast ``model_output.float()`` at entry. Passing
+        bf16 ``model_output`` must still yield an fp32 noise buffer and
+        an fp32 log-prob sum."""
+        scheduler = _DummyScheduler()
+        current_sigma = torch.tensor(0.5, dtype=torch.float32)
+        next_sigma = torch.tensor(0.3, dtype=torch.float32)
+        shape = (1, 16, 1, 32, 32)
+        pipeline_config = types.SimpleNamespace(
+            shard_latents_for_sp=lambda batch, latents: (latents, False)
+        )
+
+        for sde_type in ("sde", "cps"):
+            batch = self._build_batch(sde_type=sde_type, shape=shape)
+            scheduler.release_rollout_resources(batch)
+            scheduler.prepare_rollout(batch=batch, pipeline_config=pipeline_config)
+
+            g = torch.Generator(device="cpu").manual_seed(0)
+            model_output = torch.randn(shape, generator=g, dtype=torch.float32).to(
+                torch.bfloat16
+            )
+            sample = torch.randn(shape, generator=g, dtype=torch.float32)
+
+            # Use the real _rollout_variance_noise (no mock) so its dtype
+            # propagates from the (original) model_output.dtype into the
+            # noise buffer. If flow_sde_sampling fails to cast to fp32 at
+            # entry, the buffer is bf16 → log_prob becomes bf16.
+            scheduler.flow_sde_sampling(
+                batch,
+                model_output=model_output,
+                sample=sample,
+                current_sigma=current_sigma,
+                next_sigma=next_sigma,
+                generator=g,
+            )
+            log_prob_sum, _count = scheduler.consume_local_rollout_log_probs(batch)
+            self.assertEqual(
+                log_prob_sum.dtype,
+                torch.float32,
+                msg=f"{sde_type}: log_prob_sum must be fp32 with bf16 model_output",
+            )
+            noise_buffer = scheduler._get_rollout_session_data(batch).noise_buffer
+            self.assertEqual(
+                noise_buffer.dtype,
+                torch.float32,
+                msg=f"{sde_type}: noise_buffer must be fp32 with bf16 model_output",
+            )
+
 
 if __name__ == "__main__":
     unittest.main()