[Diffusion] Standalone Rollout API + Denoising Environment Backpass + SP-Aligned Log-Prob for T2I Post-Training

1. Architecture & Design

Motivation

Building on PR #21204 (SDE/CPS/ODE rollout log-prob engine), RL-based post-training of diffusion models (e.g., FlowGRPO) needs more than just per-step log-probabilities. It also needs (a) a dedicated serving endpoint that is insulated from the generic T2I generate path, (b) enough trajectory metadata to replay the rollout outside of SGLang-D for policy-gradient computation, and (c) numerical stability under Sequence Parallelism.

This PR delivers four things (T2I scope):

1. Standalone rollout HTTP API. Rollout is split off from the generic /generate path into its own FastAPI router POST /rollout/generate (APIRouter(prefix="/rollout")) with dedicated RolloutRequest / RolloutResponse[] I/O structs. The endpoint is sample-granular — batch results are split into per-sample RolloutResponse objects, each carrying its own serialized trajectory slice, so RL trainers consuming individual trajectories don't have to demux batches or thread RL-only fields through the main generate path.
2. Rollout denoising environment backpass. RolloutDenoisingMixin hooks DenoisingStage and, when rollout_return_denoising_env / rollout_return_dit_trajectory are set, returns the frozen transformer kwargs (text embeds, image embeds, guidance) and the per-step DiT trajectory (raw noisy x_{t_i} + final latent + timesteps, [T+1, ...]) alongside the existing log-probs / debug tensors — everything a trainer needs to re-run the exact forward pass SGLang-D ran.
3. SP-aligned log-prob. Previously rollout_log_probs under SP drifted from the single-GPU reference by exactly 1 bf16 ulp at each step's log-prob magnitude (Qwen-Image SP=2 max over 50 steps: 4.88e-4 = 2⁻¹¹ for SDE, 1.95e-3 = 2⁻⁹ for CPS). Root cause: 0-dim fp32 noise_std_dev is silently demoted to bf16 by PyTorch wrapped-scalar promotion when multiplied against an N-dim bf16 variance_noise on enable_autocast=False pipelines, and bf16 sum is non-associative across shards. Fix: each rank computes log-prob on the full pre-shard noise buffer (no all_reduce), plus flowGRPO's fp32 entry-cast policy for SDE/CPS (see §2 Precision Policy).
4. Assorted bug fixes bundled along the way — missing final-step latent in dit_trajectory, wrong capture tensor, redundant _maybe_collect_rollout_log_probs call, variance_noise aliasing on SP=1, CLI-defaults leak, gather_latents_for_sp kwarg drift. Full list in §2 "Bug Fixes".

Design Principles

• Per-sample granularity. The rollout API splits batch outputs into per-sample RolloutResponse objects, each carrying its own serialized trajectory slice, so RL trainers that process individual trajectories don't need to demux the batch.
• Opt-in fields. rollout_return_denoising_env and rollout_return_dit_trajectory are opt-in flags. When disabled, no extra memory or computation is consumed — strict zero overhead on the main path.
• SP-transparent gathering. All SP-sharded tensors (latent model inputs, cond kwargs, DiT trajectory) are gathered to full shape before being returned, using the same gather_latents_for_sp / gather_dit_env_static_for_sp infrastructure as the core pipeline. Non-rollout inference is strictly unchanged.
• Bit-exact SP log-prob, zero extra collectives. Each SP rank generates the full variance noise into rollout_session_data.noise_buffer with a common seed, then computes log_prob_no_const_val on that full buffer. Every rank gets the same per-step sum, so no all_reduce is needed — both SP/single-GPU bit-exactness and communication cost are better than the previous "local sum + all-reduce" path.
• ODE ≡ non-rollout bit-exact, preserved. The SDE/CPS branches cast model_output.float() at entry (matching flowGRPO sd3_sde_with_logprob.py), but the ODE branch is deliberately left uncasted so that rollout(sde_type="ode") produces bit-identical prev_sample to the non-rollout deterministic step in scheduling_flow_match_euler_discrete.step. A dedicated unit test locks this in with torch.equal.

Module Structure

runtime/entrypoints/post_training/
  io_struct.py        # RolloutRequest / RolloutResponse Pydantic models
  rollout_api.py      # POST /rollout/generate endpoint + batch→per-sample serialization
  utils.py            # tensor_to_base64, base64_to_tensor, _maybe_serialize, _maybe_deserialize

runtime/post_training/
  rollout_denoising_mixin.py              # DenoisingStage hooks for env/trajectory collection
  sp_utils.py                             # SP collective helpers (gather_stacked_latents_for_sp, ...)
  rl_dataclasses.py                       # RolloutTrajectoryData, RolloutDenoisingEnv, RolloutDitTrajectory
  scheduler_rl_mixin.py                   # SchedulerRLMixin — SP-aligned log-prob on full buffer
  scheduler_rl_debug_mixin.py             # SchedulerRLDebugMixin — debug tensor accumulation
  pipeline_configs/
    zimage_rollout_pipeline_mixin.py      # Z-Image cond-kwargs gather under SP
    qwen_image_rollout_pipeline_mixin.py  # Qwen-Image cond-kwargs gather under SP

test/unit/
  test_rollout_api.py                     # 31 tests: tensor serialization, response building, trajectory slicing
  test_scheduler_rollout_unit.py          #  5 tests: ODE determinism, ODE bit-exact vs non-rollout,
                                          #           flowGRPO alignment, dtype regression

Data Flow

POST /rollout/generate
  → build SamplingParams (rollout=True, rollout_return_denoising_env, rollout_return_dit_trajectory)
  → pipeline.forward(req)
      → DenoisingStage:
          _maybe_prepare_rollout(batch)
          _maybe_init_denoising_env_collection(batch, image_kwargs, pos_cond_kwargs, ...)
          for each step:
              _maybe_append_dit_trajectory_step(batch, latents, timestep)   # raw x_{t_i}, before scale
              scheduler.step(batch=batch, ...)
                → flow_sde_sampling(model_output, sample, sigma, sigma_next, generator)
                    → SDE/CPS: full-buffer log_prob on rollout_session_data.noise_buffer (fp32)
                    → ODE   : prev_sample = sample + dt * model_output (dtype-preserved)
          _postprocess_rollout_outputs
              → _maybe_collect_rollout_log_probs(batch)
              → _maybe_finalize_dit_env_collection(batch)       # appends final latent, SP-gathers env + trajectory
  → _build_response: split batch → per-sample RolloutResponse with base64-serialized tensors
  → ORJSONResponse

2. Features, API & Reliability

Rollout API Parameters (RolloutRequest + SamplingParams extensions)

Response Structure (RolloutResponse, per sample)

• generated_output: Base64-serialized output image / video tensor.
• rollout_log_probs: [T] per-step log-probability, fp32.
• rollout_debug_tensors: {variance_noise, prev_sample_mean, noise_std_dev, model_output}, each [T, ...], SP-gathered to full shape.
• denoising_env: {image_kwargs, pos_cond_kwargs, neg_cond_kwargs, guidance}, SP-gathered.
• dit_trajectory: {latents [T+1, ...], timesteps [T]}, SP-gathered.

Precision Policy (flowGRPO alignment)

The fp32 entry-cast on SDE/CPS is adopted from FlowGRPO's reference sde_step_with_logprob, whose comment reads "bf16 can overflow here when compute prev_sample_mean, we must convert all variable to fp32". This also avoids the PyTorch wrapped-scalar promotion trap where a 0-dim fp32 noise_std_dev is silently demoted to bf16 when multiplied against an N-dim bf16 variance_noise, which was the actual root cause of the pre-fix SP log-prob drift on enable_autocast=False pipelines (Qwen-Image, Z-Image).

Bug Fixes

• Missing final-step latent in dit_trajectory. _postprocess_rollout_outputs now appends the last scheduler.step output as the (T+1)-th entry and is reordered to run before _post_denoising_loop so it can be SP-gathered uniformly with the per-step trajectory.
• Wrong capture tensor for dit_trajectory. Moved the per-step capture hook to the top of the denoising loop; stores the raw pre-scale / pre-I2V-concat latent x_{t_i} instead of the post-scale latent_model_input. Field renamed latent_model_inputs → latents (shape [B, T+1, ...]).
• SP log-prob mismatch. Full-buffer log-prob computation + flowGRPO fp32 entry-cast (see §3.2 / §3.3 — was max_abs_diff=4.88e-4 SDE / 1.95e-3 CPS on Qwen-Image SP=2, now 0).
• variance_noise aliasing across debug-append calls. append_local_rollout_debug_tensors now clones before appending; on SP=1, _rollout_variance_noise returns a view of the reusable noise_buffer, so without the clone every appended entry held the last step's noise after the buffer was overwritten.
• Redundant _maybe_collect_rollout_log_probs call. Removed the duplicate call from _post_denoising_loop — it was already invoked from _postprocess_rollout_outputs, and the first one crashed on SP runs because _rollout_session_data had already been released.
• Rollout CLI defaults leaking into sampling_params. Removed explicit defaults from rollout CLI args so test_get_cli_args_drops_unset_sampling_params passes again.
• gather_latents_for_sp keyword-arg drift in Z-Image / Qwen-Image rollout pipeline mixins and the debug mixin — small positional/keyword arg fix.

Reliability Testing

Three levels of testing validate correctness:

1. Unit tests — the full CI unit suite (python3 sglang/multimodal_gen/test/run_suite.py --suite unit) runs 125 tests (124 passed, 1 env-dependent skip: test_integration_with_moto). Of those, 36 are rollout-specific (31 in test_rollout_api.py + 5 in test_scheduler_rollout_unit.py); see §3.1.
2. FlowGRPO reference alignment: SDE/CPS single-step computation is compared against a verbatim copy of FlowGRPO's sde_step_with_logprob across 4 seeds, for all four output quantities (prev_sample, prev_sample_mean, noise_std_dev, log_prob). All < 1e-6.
3. Multi-parallel-config trajectory debug on Qwen-Image (50-step, 1024×1024, 2 GPUs) under TP1-SP2 / TP2-SP1 / TP1-SP1-CFGP, comparing all debug tensors + log-prob against the single-GPU reference — see §3.2. A secondary Z-Image-Turbo run with the same harness is in Appendix A.

3. Detailed Test Results

3.1 Unit Tests — Full CI unit suite 124 passed / 1 skipped

Command (matches .github/workflows/pr-test-multimodal-gen.yml::multimodal-gen-unit-test):

cd python && python3 sglang/multimodal_gen/test/run_suite.py --suite unit
# => 124 passed, 1 skipped, 3 warnings, 8 subtests passed in 9.91s
# (skipped = test_storage.py::test_integration_with_moto, env-dependent S3-moto integration)

Rollout-specific breakdown (36/36):

python/sglang/multimodal_gen/test/unit/test_rollout_api.py             31 passed
python/sglang/multimodal_gen/test/unit/test_scheduler_rollout_unit.py   5 passed

The 5 scheduler rollout tests lock in the following invariants:

3.2 Qwen-Image Parallel Consistency (primary)

Model: Qwen/Qwen-Image, seed=42, noise_level=0.5, 50 steps, 1024×1024, 2 GPUs.
Reference: single-GPU (TP1, SP1, no CFGP).
Prompt: "A fluffy silver-gray cat rests on a light surface indoors, wearing small round sunglasses with gold frames and dark blue lenses. …"

SDE mode

CPS mode

ODE mode

3.3 log_prob before vs after the SP fix

Historical Qwen-Image TP1-SP2 run (pre-fix, from outputs/rollout_trajectory_debug_compare/20260411-080016/trajectory_debug_report.md):

Key observations (Qwen-Image, primary)

• Qwen-Image SP2 is bit-exact on every tensor including log_prob.
• Qwen-Image TP2-SP1 is also bit-exact across all 5 metrics.
• CFGP (tp1_sp1_cfg1) shows small prev_sample_mean and model_output differences from non-deterministic reduction order in the CFG-parallel transformer forward — not the rollout engine. log_prob, variance_noise and noise_std_dev are all bit-exact.
• ODE rollout ↔ non-rollout deterministic step is bit-exact at the step-function level (guarded by test_ode_bit_exact_with_non_rollout_path). Any residual drift under TP2/CFGP is exactly what the non-rollout path would also see, consistent with PR [Diffusion] Revamp Rollout Log-Prob Support with SDE/CPS for RL Post-Training #21204's baseline.
• Appendix A replicates the same experiment on Z-Image-Turbo and shows the same pattern (SP2 bit-exact, log-prob bit-exact on every config; slightly more DiT-side FP drift under TP2-SP1 because of a different DiT architecture).

4. Summary

• Standalone POST /rollout/generate API (new FastAPI router, separate from /generate), sample-granular, returns per-sample serialized rollout trajectory data (log-probs, debug tensors, denoising env, DiT trajectory) for T2I pipelines.
• Denoising-environment backpass: RolloutDenoisingMixin captures frozen transformer kwargs + per-step x_{t_i} + final latent, SP-gathered for replay.
• SP-aligned rollout_log_probs, bit-exact with the single-GPU reference on Qwen-Image and Z-Image-Turbo across all three modes (sde / cps / ode). Achieved by computing log-prob on the full pre-shard noise buffer + flowGRPO's fp32 entry-cast policy for SDE/CPS, while ODE preserves bit-exactness with the non-rollout deterministic path (locked in by unit test).
• Assorted bug fixes — see §2 "Bug Fixes" for the full list.
• Full CI unit suite (run_suite.py --suite unit) 124 passed / 1 env-skipped, including the 36 rollout-specific tests. Two new invariants lock in (a) ODE ≡ non-rollout bit-exactness and (b) the fp32 entry-cast regression guard against PyTorch wrapped-scalar promotion. SDE/CPS single-step also matches a verbatim FlowGRPO reference across 4 seeds within 1e-6.
• Total python diff against upstream/main: 20 files, +1276 / −78 lines (excluding unrelated upstream changes).

Appendix A — Z-Image-Turbo parallel consistency (secondary verification)

Z-Image-Turbo was run through the same test_rollout_trajectory_debug_parallel.py harness as Qwen-Image to make sure the SP fix generalizes beyond a single model. Same prompt/seed/noise_level as §3.2; 9 steps instead of 50 because Z-Image-Turbo is turbo-distilled.

Model: Tongyi-MAI/Z-Image-Turbo, seed=42, noise_level=0.5, 9 steps, 1024×1024, 2 GPUs.
Reference: single-GPU (TP1, SP1, no CFGP). max_abs_diff is the max over all 9 steps.

SDE mode

CPS mode

ODE mode

Z-Image-Turbo observation: SP2 is bit-exact on every tensor including log_prob. TP2-SP1 shows slightly more prev_sample_mean / model_output drift than Qwen-Image does (reduction-order non-determinism in a different DiT architecture), but log_prob, variance_noise and noise_std_dev are all still bit-exact, confirming the full-buffer log-prob path is insulated from downstream DiT non-determinism regardless of the backbone.