[diffusion] feat: add rollout log_prob with flow-matching SDE/CPS support

Rebased onto latest main.

Author: zhaochenyang20

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

@@ -159,6 +159,9 @@ class SamplingParams:
     # Misc
     save_output: bool = True
     return_frames: bool = False
+    rollout: bool = False
+    rollout_sde_type: str = "sde"
+    rollout_noise_level: float = 0.7
     return_trajectory_latents: bool = False  # returns all latents for each timestep
     return_trajectory_decoded: bool = False  # returns decoded latents for each timestep
     # if True, disallow user params to override subclass-defined protected fields
@@ -306,6 +309,9 @@ def _finite_non_negative_float(
         _finite_non_negative_float(
             "guidance_rescale", self.guidance_rescale, allow_none=False
         )
+        _finite_non_negative_float(
+            "rollout_noise_level", self.rollout_noise_level, allow_none=False
+        )
 
         if self.cfg_normalization is None:
             self.cfg_normalization = 0.0
@@ -808,6 +814,25 @@ def add_cli_args(parser: Any) -> Any:
             default=SamplingParams.return_trajectory_latents,
             help="Whether to return the trajectory",
         )
+        parser.add_argument(
+            "--rollout",
+            action="store_true",
+            default=SamplingParams.rollout,
+            help="Enable rollout mode and return per-step log_prob trajectory",
+        )
+        parser.add_argument(
+            "--rollout-sde-type",
+            type=str,
+            choices=["sde", "cps"],
+            default=SamplingParams.rollout_sde_type,
+            help="Rollout step objective type used in log-prob computation.",
+        )
+        parser.add_argument(
+            "--rollout-noise-level",
+            type=float,
+            default=SamplingParams.rollout_noise_level,
+            help="Noise level used by rollout SDE/CPS step objective.",
+        )
         parser.add_argument(
             "--return-trajectory-decoded",
             action="store_true",

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

@@ -216,6 +216,7 @@ def generate(
                         ),
                         trajectory_latents=output_batch.trajectory_latents,
                         trajectory_timesteps=output_batch.trajectory_timesteps,
+                        trajectory_log_probs=output_batch.trajectory_log_probs,
                         trajectory_decoded=output_batch.trajectory_decoded,
                     )
 

python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py

@@ -129,6 +129,9 @@ async def generations(
             true_cfg_scale=request.true_cfg_scale,
             negative_prompt=request.negative_prompt,
             enable_teacache=request.enable_teacache,
+            rollout=request.rollout,
+            rollout_sde_type=request.rollout_sde_type,
+            rollout_noise_level=request.rollout_noise_level,
             output_compression=request.output_compression,
             output_quality=request.output_quality,
         )

python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py

@@ -46,6 +46,9 @@ class ImageGenerationsRequest(BaseModel):
     output_quality: Optional[str] = "default"
     output_compression: Optional[int] = None
     enable_teacache: Optional[bool] = False
+    rollout: Optional[bool] = False
+    rollout_sde_type: Optional[str] = "sde"
+    rollout_noise_level: Optional[float] = 0.7
     diffusers_kwargs: Optional[Dict[str, Any]] = None  # kwargs for diffusers backend
 
 
@@ -98,6 +101,9 @@ class VideoGenerationsRequest(BaseModel):
     output_quality: Optional[str] = "default"
     output_compression: Optional[int] = None
     output_path: Optional[str] = None
+    rollout: Optional[bool] = False
+    rollout_sde_type: Optional[str] = "sde"
+    rollout_noise_level: Optional[float] = 0.7
     diffusers_kwargs: Optional[Dict[str, Any]] = None  # kwargs for diffusers backend
 
 

python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py

@@ -75,6 +75,9 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque
         frame_interpolation_exp=request.frame_interpolation_exp,
         frame_interpolation_scale=request.frame_interpolation_scale,
         frame_interpolation_model_path=request.frame_interpolation_model_path,
+        rollout=request.rollout,
+        rollout_sde_type=request.rollout_sde_type,
+        rollout_noise_level=request.rollout_noise_level,
         output_path=request.output_path,
         output_compression=request.output_compression,
         output_quality=request.output_quality,
@@ -181,6 +184,9 @@ async def create_video(
     frame_interpolation_exp: Optional[int] = Form(1),
     frame_interpolation_scale: Optional[float] = Form(1.0),
     frame_interpolation_model_path: Optional[str] = Form(None),
+    rollout: Optional[bool] = Form(False),
+    rollout_sde_type: Optional[str] = Form("sde"),
+    rollout_noise_level: Optional[float] = Form(0.7),
     output_quality: Optional[str] = Form("default"),
     output_compression: Optional[int] = Form(None),
     extra_body: Optional[str] = Form(None),
@@ -256,6 +262,9 @@ async def create_video(
             frame_interpolation_exp=frame_interpolation_exp,
             frame_interpolation_scale=frame_interpolation_scale,
             frame_interpolation_model_path=frame_interpolation_model_path,
+            rollout=rollout,
+            rollout_sde_type=rollout_sde_type,
+            rollout_noise_level=rollout_noise_level,
             output_compression=output_compression,
             output_quality=output_quality,
             **(

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

@@ -108,6 +108,7 @@ class GenerationResult:
     metrics: dict = field(default_factory=dict)
     trajectory_latents: Any = None
     trajectory_timesteps: Any = None
+    trajectory_log_probs: Any = None
     trajectory_decoded: Any = None
     prompt_index: int = 0
     output_file_path: str | None = None

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -233,6 +233,7 @@ def execute_forward(self, batch: List[Req]) -> OutputBatch:
                     metrics=result.metrics,
                     trajectory_timesteps=getattr(result, "trajectory_timesteps", None),
                     trajectory_latents=getattr(result, "trajectory_latents", None),
+                    trajectory_log_probs=getattr(result, "trajectory_log_probs", None),
                     noise_pred=getattr(result, "noise_pred", None),
                     trajectory_decoded=getattr(result, "trajectory_decoded", None),
                 )

python/sglang/multimodal_gen/runtime/pipelines/patches/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

python/sglang/multimodal_gen/runtime/pipelines/patches/flow_matching_with_logprob.py

@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Flow-matching rollout step utilities for log-prob computation."""
+
+import math
+from typing import Any, Optional, Union
+
+import torch
+from diffusers.utils.torch_utils import randn_tensor
+
+
+def sde_step_with_logprob(
+    self: Any,
+    model_output: torch.FloatTensor,
+    sample: torch.FloatTensor,
+    step_index: int,
+    noise_level: float = 0.7,
+    prev_sample: Optional[torch.FloatTensor] = None,
+    generator: Optional[Union[torch.Generator, list[torch.Generator]]] = None,
+    sde_type: str = "sde",
+):
+    """Run one rollout step and compute per-sample log_prob.
+
+    sde_type="sde" uses the Gaussian transition objective.
+    sde_type="cps" uses the simplified CPS objective.
+    """
+    sample_dtype = sample.dtype
+    model_output = model_output.float()
+    sample = sample.float()
+    if prev_sample is not None:
+        prev_sample = prev_sample.float()
+
+    step_indices = torch.full(
+        (sample.shape[0],),
+        int(step_index),
+        device=self.sigmas.device,
+        dtype=torch.long,
+    )
+    prev_step_indices = (step_indices + 1).clamp_max(len(self.sigmas) - 1)
+    sigma = self.sigmas[step_indices].to(device=sample.device, dtype=sample.dtype)
+    sigma_prev = self.sigmas[prev_step_indices].to(
+        device=sample.device, dtype=sample.dtype
+    )
+    sigma = sigma.view(-1, *([1] * (sample.ndim - 1)))
+    sigma_prev = sigma_prev.view(-1, *([1] * (sample.ndim - 1)))
+    sigma_max = self.sigmas[min(1, len(self.sigmas) - 1)].to(
+        device=sample.device, dtype=sample.dtype
+    )
+    dt = sigma_prev - sigma
+
+    if sde_type == "sde":
+        denom_sigma = 1 - torch.where(
+            torch.isclose(sigma, sigma.new_tensor(1.0)), sigma_max, sigma
+        )
+        std_dev_t = torch.sqrt((sigma / denom_sigma).clamp_min(1e-12)) * noise_level
+        prev_sample_mean = (
+            sample * (1 + std_dev_t**2 / (2 * sigma) * dt)
+            + model_output * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) * dt
+        )
+
+        sqrt_neg_dt = torch.sqrt((-dt).clamp_min(1e-12))
+        if prev_sample is None:
+            variance_noise = randn_tensor(
+                model_output.shape,
+                generator=generator,
+                device=model_output.device,
+                dtype=model_output.dtype,
+            )
+            prev_sample = prev_sample_mean + std_dev_t * sqrt_neg_dt * variance_noise
+
+        std = (std_dev_t * sqrt_neg_dt).clamp_min(1e-12)
+        log_prob = (
+            -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * (std**2))
+            - torch.log(std)
+            - torch.log(torch.sqrt(torch.as_tensor(2 * math.pi, device=std.device)))
+        )
+    elif sde_type == "cps":
+        std_dev_t = sigma_prev * math.sin(noise_level * math.pi / 2)
+        pred_original_sample = sample - sigma * model_output
+        noise_estimate = sample + model_output * (1 - sigma)
+        sigma_delta = (sigma_prev**2 - std_dev_t**2).clamp_min(0.0)
+        prev_sample_mean = pred_original_sample * (
+            1 - sigma_prev
+        ) + noise_estimate * torch.sqrt(sigma_delta)
+
+        if prev_sample is None:
+            variance_noise = randn_tensor(
+                model_output.shape,
+                generator=generator,
+                device=model_output.device,
+                dtype=model_output.dtype,
+            )
+            prev_sample = prev_sample_mean + std_dev_t * variance_noise
+
+        # CPS transition is Gaussian with std_dev_t, so compute a valid log-probability.
+        std = std_dev_t.clamp_min(1e-12)
+        log_prob = (
+            -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * (std**2))
+            - torch.log(std)
+            - torch.log(torch.sqrt(torch.as_tensor(2 * math.pi, device=std.device)))
+        )
+    else:
+        raise ValueError(f"Unsupported sde_type: {sde_type}")
+
+    log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim)))
+    return prev_sample.to(sample_dtype), log_prob

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

@@ -133,8 +133,9 @@ class Req:
     # Component modules (populated by the pipeline)
     modules: dict[str, Any] = field(default_factory=dict)
 
-    trajectory_timesteps: list[torch.Tensor] | None = None
+    trajectory_timesteps: torch.Tensor | None = None
     trajectory_latents: torch.Tensor | None = None
+    trajectory_log_probs: torch.Tensor | None = None
     trajectory_audio_latents: torch.Tensor | None = None
 
     # Extra parameters that might be needed by specific pipeline implementations
@@ -329,8 +330,9 @@ class OutputBatch:
     output: torch.Tensor | None = None
     audio: torch.Tensor | None = None
     audio_sample_rate: int | None = None
-    trajectory_timesteps: list[torch.Tensor] | None = None
+    trajectory_timesteps: torch.Tensor | None = None
     trajectory_latents: torch.Tensor | None = None
+    trajectory_log_probs: torch.Tensor | None = None
     trajectory_decoded: list[torch.Tensor] | None = None
     error: str | None = None
     output_file_paths: list[str] | None = None