chore: import upstream snapshot with attribution

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

@@ -0,0 +1,110 @@
+"""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
+class UpdateWeightFromDiskReqInput:
+    """Request to update model weights from disk for diffusion models."""
+
+    model_path: str
+    flush_cache: bool = True
+    target_modules: list[str] | None = None
+
+
+@dataclass
+class UpdateWeightFromTensorReqInput:
+    """Request to update model weights from tensor payloads for diffusion models."""
+
+    serialized_named_tensors: list[str | bytes]
+    load_format: str | None = None
+    target_modules: list[str] | None = None
+
+
+@dataclass
+class UpdateWeightFromTensorCheckerReqInput:
+    """Request to verify live module weights against expected SHA-256 values."""
+
+    target_module: str
+    expected_named_tensors_sha256: dict[str, str]
+
+
+@dataclass
+class GetWeightsChecksumReqInput:
+    """Compute SHA-256 checksum of loaded module weights for verification."""
+
+    module_names: list[str] | None = None
+
+
+@dataclass
+class ReleaseMemoryOccupationReqInput:
+    """Request to release (sleep) GPU memory occupation for the diffusion engine."""
+
+    pass
+
+
+@dataclass
+class ResumeMemoryOccupationReqInput:
+    """Request to resume (wake) GPU memory occupation for the diffusion engine."""
+
+    pass
+
+
+class RolloutRequest(BaseModel):
+    prompt: str
+    negative_prompt: Optional[str] = None
+    seed: Optional[int] = None
+    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
+
+    # 0-indexed denoising-loop step filters. None = all steps.
+    rollout_sde_step_indices: Optional[list[int]] = None
+    rollout_return_step_indices: Optional[list[int]] = None
+
+    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,329 @@
+"""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, *, current_key: str | None = None
+) -> 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, current_key=k)
+            for k, v in obj.items()
+        }
+    if isinstance(obj, list):
+        if current_key == "img_shapes" and len(obj) == batch_size:
+            return [obj[sample_idx]]
+        return [
+            _extract_single_sample_tensor(
+                v, sample_idx, batch_size, current_key=current_key
+            )
+            for v in obj
+        ]
+    if isinstance(obj, tuple):
+        return tuple(
+            _extract_single_sample_tensor(
+                v, sample_idx, batch_size, current_key=current_key
+            )
+            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]
+        if isinstance(out_i, torch.Tensor):
+            out_i = out_i.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
+
+
+def _build_sampling_kwargs(request: RolloutRequest) -> dict:
+    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,
+        rollout_sde_step_indices=request.rollout_sde_step_indices,
+        rollout_return_step_indices=request.rollout_return_step_indices,
+        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
+    return {k: v for k, v in sampling_kwargs.items() if v is not None}
+
+
+@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 = _build_sampling_kwargs(request)
+    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,48 @@
+"""Tensor serialization for post-training / rollout HTTP responses."""
+
+from __future__ import annotations
+
+import base64
+from typing import Any
+
+import numpy as np
+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, np.ndarray):
+        return _maybe_serialize(torch.from_numpy(obj))
+    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/entrypoints/post_training/weights_api.py

@@ -0,0 +1,199 @@
+"""Weight update API for the diffusion engine."""
+
+from fastapi import APIRouter, Request
+
+from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import (
+    GetWeightsChecksumReqInput,
+    ReleaseMemoryOccupationReqInput,
+    ResumeMemoryOccupationReqInput,
+    UpdateWeightFromDiskReqInput,
+    UpdateWeightFromTensorCheckerReqInput,
+    UpdateWeightFromTensorReqInput,
+)
+from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
+from sglang.srt.utils.json_response import orjson_response
+
+router = APIRouter()
+
+
+@router.post("/update_weights_from_disk")
+async def update_weights_from_disk(request: Request):
+    """Update model weights from disk inplace without restarting the server."""
+    body = await request.json()
+    model_path = body.get("model_path")
+    if not model_path:
+        return orjson_response(
+            {"success": False, "message": "model_path is required"},
+            status_code=400,
+        )
+
+    req = UpdateWeightFromDiskReqInput(
+        model_path=model_path,
+        flush_cache=body.get("flush_cache", True),
+        target_modules=body.get("target_modules"),
+    )
+
+    try:
+        response = await async_scheduler_client.forward(req)
+    except Exception as e:
+        return orjson_response(
+            {"success": False, "message": str(e)},
+            status_code=500,
+        )
+
+    if response.output is None:
+        return orjson_response(
+            {
+                "success": False,
+                "message": response.error or "Unknown status",
+            },
+            status_code=500,
+        )
+
+    result = response.output
+    return orjson_response(
+        result,
+        status_code=200 if result["success"] else 400,
+    )
+
+
+@router.post("/update_weights_from_tensor")
+async def update_weights_from_tensor(request: Request):
+    """Update model weights from serialized tensor payloads."""
+    body = await request.json()
+    serialized_named_tensors = body.get("serialized_named_tensors")
+    if not serialized_named_tensors:
+        return orjson_response(
+            {"success": False, "message": "serialized_named_tensors is required"},
+            status_code=400,
+        )
+
+    req = UpdateWeightFromTensorReqInput(
+        serialized_named_tensors=serialized_named_tensors,
+        load_format=body.get("load_format"),
+        target_modules=body.get("target_modules"),
+    )
+
+    try:
+        response = await async_scheduler_client.forward(req)
+    except Exception as e:
+        return orjson_response(
+            {"success": False, "message": str(e)},
+            status_code=500,
+        )
+
+    result = response.output
+    return orjson_response(
+        result,
+        status_code=200 if result["success"] else 400,
+    )
+
+
+@router.post("/update_weights_from_tensor_checker")
+async def update_weights_from_tensor_checker(request: Request):
+    """Verify live module weights against expected SHA-256 values."""
+    body = await request.json()
+    target_module = body.get("target_module")
+    if not target_module:
+        return orjson_response(
+            {"success": False, "message": "target_module is required"},
+            status_code=400,
+        )
+
+    expected_named_tensors_sha256 = body.get("expected_named_tensors_sha256")
+    if (
+        not isinstance(expected_named_tensors_sha256, dict)
+        or not expected_named_tensors_sha256
+    ):
+        return orjson_response(
+            {
+                "success": False,
+                "message": "expected_named_tensors_sha256 is required",
+            },
+            status_code=400,
+        )
+
+    req = UpdateWeightFromTensorCheckerReqInput(
+        target_module=target_module,
+        expected_named_tensors_sha256=expected_named_tensors_sha256,
+    )
+
+    try:
+        response = await async_scheduler_client.forward(req)
+    except Exception as e:
+        return orjson_response(
+            {"success": False, "message": str(e)},
+            status_code=500,
+        )
+
+    result = response.output
+    success = result.get("success", False)
+    message = result.get("message", "Unknown status")
+    return orjson_response(
+        {"success": success, "message": message},
+        status_code=200 if success else 400,
+    )
+
+
+@router.post("/get_weights_checksum")
+async def get_weights_checksum(request: Request):
+    """Return SHA-256 checksum of each requested module's weights."""
+    body = await request.json()
+    req = GetWeightsChecksumReqInput(
+        module_names=body.get("module_names"),
+    )
+
+    try:
+        response = await async_scheduler_client.forward(req)
+    except Exception as e:
+        return orjson_response({"error": str(e)}, status_code=500)
+
+    return orjson_response(response.output, status_code=200)
+
+
+@router.post("/release_memory_occupation")
+async def release_memory_occupation():
+    """Release GPU memory occupation (sleep the engine)."""
+    try:
+        response = await async_scheduler_client.forward(
+            ReleaseMemoryOccupationReqInput()
+        )
+    except Exception as e:
+        return orjson_response({"success": False, "message": str(e)}, status_code=500)
+
+    if response.output is None:
+        return orjson_response(
+            {
+                "success": False,
+                "message": response.error or "Unknown status",
+            },
+            status_code=500,
+        )
+
+    payload = response.output
+    success = bool(payload["success"])
+    return orjson_response(payload, status_code=200 if success else 400)
+
+
+@router.post("/resume_memory_occupation")
+async def resume_memory_occupation():
+    """Resume GPU memory occupation (wake the engine)."""
+    try:
+        response = await async_scheduler_client.forward(
+            ResumeMemoryOccupationReqInput()
+        )
+    except Exception as e:
+        return orjson_response({"success": False, "message": str(e)}, status_code=500)
+
+    if response.output is None:
+        return orjson_response(
+            {
+                "success": False,
+                "message": response.error or "Unknown status",
+            },
+            status_code=500,
+        )
+
+    payload = response.output
+    success = bool(payload["success"])
+    return orjson_response(payload, status_code=200 if success else 400)