local: merge PR sgl-project#18806 (rollout log_prob) + PR sgl-project#19153 (sleep/wake)

Cherry-picked from:
- PR sgl-project#18806 (MikukuOvO): flow-matching SDE/CPS log_prob
- PR sgl-project#19153 (Godmook): release/resume memory occupation

Known issues:
- t.item() in log_prob path causes GPU sync overhead
- release_memory_occupation tags only supports "weights"

Author: haonan3

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

@@ -146,6 +146,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
@@ -293,6 +296,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
@@ -743,6 +749,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

@@ -18,6 +18,8 @@
     GenerationResult,
     ListLorasReq,
     MergeLoraWeightsReq,
+    ReleaseMemoryOccupationReq,
+    ResumeMemoryOccupationReq,
     SetLoraReq,
     ShutdownReq,
     UnmergeLoraWeightsReq,
@@ -213,6 +215,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,
                     )
 
@@ -452,6 +455,40 @@ def generate_with_lora(
             )
         )
 
+    def release_memory_occupation(self, tags: List[str] | None = None) -> dict:
+        """Release GPU memory (sleep). Offloads model weights to CPU.
+
+        Args:
+            tags: Which memory regions to release. Currently only "weights" is
+                  supported for diffusion. If omitted, all regions are released.
+
+        Returns:
+            dict with "success" and "message" keys.
+        """
+        req = ReleaseMemoryOccupationReq(tags=tags)
+        response = sync_scheduler_client.forward(req)
+        if response.error:
+            raise RuntimeError(f"Failed to release memory: {response.error}")
+        logger.info("Successfully released GPU memory occupation (sleeping).")
+        return response.output
+
+    def resume_memory_occupation(self, tags: List[str] | None = None) -> dict:
+        """Resume GPU memory (wake up). Loads model weights back to GPU.
+
+        Args:
+            tags: Which memory regions to resume. Currently only "weights" is
+                  supported for diffusion. If omitted, all regions are resumed.
+
+        Returns:
+            dict with "success" and "message" keys.
+        """
+        req = ResumeMemoryOccupationReq(tags=tags)
+        response = sync_scheduler_client.forward(req)
+        if response.error:
+            raise RuntimeError(f"Failed to resume memory: {response.error}")
+        logger.info("Successfully resumed GPU memory occupation (waking up).")
+        return response.output
+
     def shutdown(self):
         """
         Shutdown the generator.

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

@@ -19,6 +19,8 @@
 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.utils import (
+    ReleaseMemoryOccupationReq,
+    ResumeMemoryOccupationReq,
     prepare_request,
     save_outputs,
 )
@@ -101,6 +103,78 @@ async def health_generate():
     return {"status": "ok"}
 
 
+def _extract_tags_from_body(body: dict) -> list[str] | None:
+    """Return the ``tags`` field from a parsed request body, or ``None``."""
+    if not isinstance(body, dict):
+        return None
+    tags = body.get("tags", None)
+    if tags is not None and not isinstance(tags, list):
+        raise ValueError(
+            f"'tags' must be a list of strings, got: {type(tags).__name__}"
+        )
+    return tags
+
+
+@health_router.post("/release_memory_occupation")
+async def release_memory_occupation(request: Request):
+    """Release GPU memory occupation (sleep).
+
+    Offloads all model weights to CPU so the GPU is free for another
+    workload (e.g. RL training).  The server process stays alive;
+    call ``/resume_memory_occupation`` to reload weights before the
+    next generation.
+
+    Body (optional JSON):
+        tags (list[str]): memory regions to release.
+            Supported value: ``"weights"``.  Omit to release all.
+    """
+    try:
+        body = await request.json()
+    except Exception:
+        body = {}
+    try:
+        tags = _extract_tags_from_body(body)
+    except ValueError as exc:
+        return ORJSONResponse({"success": False, "message": str(exc)}, status_code=422)
+
+    req = ReleaseMemoryOccupationReq(tags=tags)
+    response = await async_scheduler_client.forward(req)
+    if response.error:
+        return ORJSONResponse(
+            {"success": False, "message": response.error}, status_code=400
+        )
+    return ORJSONResponse(response.output)
+
+
+@health_router.post("/resume_memory_occupation")
+async def resume_memory_occupation(request: Request):
+    """Resume GPU memory occupation (wake up).
+
+    Loads model weights back onto the GPU so the server can serve
+    generation requests again.
+
+    Body (optional JSON):
+        tags (list[str]): memory regions to resume.
+            Supported value: ``"weights"``.  Omit to resume all.
+    """
+    try:
+        body = await request.json()
+    except Exception:
+        body = {}
+    try:
+        tags = _extract_tags_from_body(body)
+    except ValueError as exc:
+        return ORJSONResponse({"success": False, "message": str(exc)}, status_code=422)
+
+    req = ResumeMemoryOccupationReq(tags=tags)
+    response = await async_scheduler_client.forward(req)
+    if response.error:
+        return ORJSONResponse(
+            {"success": False, "message": response.error}, status_code=400
+        )
+    return ORJSONResponse(response.output)
+
+
 def make_serializable(obj):
     """Recursively converts Tensors to None for JSON serialization."""
     if isinstance(obj, torch.Tensor):

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

@@ -117,6 +117,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
 
 
@@ -93,6 +96,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

@@ -69,6 +69,9 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque
         guidance_scale_2=request.guidance_scale_2,
         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_path=request.output_path,
         output_compression=request.output_compression,
         output_quality=request.output_quality,
@@ -159,6 +162,9 @@ async def create_video(
     guidance_scale: Optional[float] = Form(None),
     num_inference_steps: Optional[int] = Form(None),
     enable_teacache: Optional[bool] = Form(False),
+    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),
@@ -212,6 +218,9 @@ async def create_video(
             negative_prompt=negative_prompt,
             num_inference_steps=num_inference_steps,
             enable_teacache=enable_teacache,
+            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

@@ -69,6 +69,20 @@ class ShutdownReq:
     pass
 
 
+@dataclass
+class ReleaseMemoryOccupationReq:
+    """Request to release GPU memory (sleep). Offloads model weights to CPU."""
+
+    tags: Optional[List[str]] = None
+
+
+@dataclass
+class ResumeMemoryOccupationReq:
+    """Request to resume GPU memory (wake up). Loads model weights back to GPU."""
+
+    tags: Optional[List[str]] = None
+
+
 def format_lora_message(
     lora_nickname: Union[str, List[str]],
     target: Union[str, List[str]],
@@ -108,6 +122,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