[diffusion] fix: address rollout log_prob review feedback (#18806)

Author: MikukuOvO

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

@@ -53,15 +53,14 @@ def sde_step_with_logprob(
     timestep_tensor = _as_timestep_tensor(timestep, batch_size, sample.device)
     step_indices = torch.tensor(
         [self.index_for_timestep(t.to(self.timesteps.device)) for t in timestep_tensor],
-        device=sample.device,
+        device=self.sigmas.device,
         dtype=torch.long,
     )
     prev_step_indices = (step_indices + 1).clamp_max(len(self.sigmas) - 1)
-    step_indices = step_indices.to(device=self.sigmas.device)
-    prev_step_indices = prev_step_indices.to(device=self.sigmas.device)
-
-    sigma = self.sigmas[step_indices].to(sample.device).to(sample.dtype)
-    sigma_prev = self.sigmas[prev_step_indices].to(sample.device).to(sample.dtype)
+    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(
@@ -70,7 +69,9 @@ def sde_step_with_logprob(
     dt = sigma_prev - sigma
 
     if sde_type == "sde":
-        denom_sigma = 1 - torch.where(sigma == 1, sigma_max, sigma)
+        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)
@@ -117,4 +118,4 @@ def sde_step_with_logprob(
         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, prev_sample_mean, std_dev_t
+    return prev_sample.to(sample_dtype), log_prob

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

@@ -113,9 +113,6 @@ def __init__(
             self._maybe_enable_torch_compile(transformer)
 
         self.scheduler = scheduler
-        self.scheduler.sde_step_with_logprob = sde_step_with_logprob.__get__(
-            self.scheduler, type(self.scheduler)
-        )
         self.vae = vae
         self.pipeline = weakref.ref(pipeline) if pipeline else None
 
@@ -1014,27 +1011,9 @@ def forward(
         trajectory_log_probs: list[torch.Tensor] = []
         rollout_enabled = bool(batch.rollout)
         rollout_sde_type = batch.rollout_sde_type
-        if rollout_sde_type is None or str(rollout_sde_type).strip() == "":
-            if rollout_enabled:
-                logger.warning("rollout_sde_type is not set, defaulting to 'sde'.")
-            rollout_sde_type = "sde"
-        else:
-            rollout_sde_type = str(rollout_sde_type).strip().lower()
-        if rollout_sde_type not in ("sde", "cps"):
-            logger.warning(
-                "Unknown rollout_sde_type '%s', using default 'sde'.",
-                rollout_sde_type,
-            )
-            rollout_sde_type = "sde"
 
         rollout_noise_level = batch.rollout_noise_level
 
-        if rollout_enabled and not hasattr(self.scheduler, "sde_step_with_logprob"):
-            raise RuntimeError(
-                f"Rollout is enabled, but scheduler '{type(self.scheduler).__name__}' "
-                "does not provide sde_step_with_logprob."
-            )
-
         # Run denoising loop
         denoising_start_time = time.time()
 
@@ -1122,15 +1101,14 @@ def forward(
 
                         # Compute the previous noisy sample
                         if rollout_enabled:
-                            latents, step_log_prob, _, _ = (
-                                self.scheduler.sde_step_with_logprob(
-                                    model_output=noise_pred,
-                                    timestep=t_device,
-                                    sample=latents,
-                                    generator=batch.generator,
-                                    sde_type=rollout_sde_type,
-                                    noise_level=rollout_noise_level,
-                                )
+                            latents, step_log_prob = sde_step_with_logprob(
+                                self.scheduler,
+                                model_output=noise_pred,
+                                timestep=t_device,
+                                sample=latents,
+                                generator=batch.generator,
+                                sde_type=rollout_sde_type,
+                                noise_level=rollout_noise_level,
                             )
                             trajectory_log_probs.append(step_log_prob)
                         else: