[diffusion] feat: opt-in uint8 rollout video via rollout_video_dtype

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

@@ -83,6 +83,9 @@ class RolloutRequest(BaseModel):
     rollout_noise_level: float = 0.7
     rollout_log_prob_no_const: bool = False
     rollout_debug_mode: bool = True
+    # "keep" returns the decoded video as-is; "uint8" quantises it engine-side
+    # with the reward path's own formula, cutting the response body ~4x.
+    rollout_video_dtype: str = "keep"
 
     rollout_return_denoising_env: bool = False
     rollout_return_dit_trajectory: bool = False

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

@@ -200,8 +200,21 @@ def _serialize_rollout_trajectory(
     )
 
 
+def _quantize_video_uint8(video: torch.Tensor) -> torch.Tensor:
+    """Same formula the training-side reward path applies before scoring."""
+    out = video.float()
+    if float(out.max()) <= 1.0 + 1e-3:
+        out = out * 255.0
+    return out.clamp(0, 255).to(torch.uint8)
+
+
 def _build_response(
-    request_id: str, prompt: str, seed: int, rollout: bool, result: OutputBatch
+    request_id: str,
+    prompt: str,
+    seed: int,
+    rollout: bool,
+    result: OutputBatch,
+    video_dtype: str = "keep",
 ) -> list[RolloutResponse]:
     """
     rollout: bool - set to False when evaluating the model
@@ -233,6 +246,8 @@ def _build_response(
     for sample_idx in range(batch_size):
         out_i = result.output[sample_idx]
         if isinstance(out_i, torch.Tensor):
+            if video_dtype == "uint8" and out_i.is_floating_point():
+                out_i = _quantize_video_uint8(out_i)
             out_i = out_i.contiguous()
         serialized_generated_output = _maybe_serialize(out_i)
         if not rollout:
@@ -353,7 +368,12 @@ async def rollout_generate(request: RolloutRequest):
         raise HTTPException(status_code=500, detail=output_batch.error)
     stamps.mark("build_start")
     rollout_responses = _build_response(
-        request_id, request.prompt, request.seed, request.rollout, output_batch
+        request_id,
+        request.prompt,
+        request.seed,
+        request.rollout,
+        output_batch,
+        video_dtype=request.rollout_video_dtype,
     )
     stamps.mark("build_end")
     payload = [r.model_dump() for r in rollout_responses]