[Diffusion] Revamp Rollout Log-Prob Support with SDE/CPS for RL Post-Training

1. Architecture & Design

Motivation

RL-based post-training of diffusion models (e.g., FlowGRPO) requires computing per-step log-probabilities along the denoising trajectory. This PR adds a modular rollout log-prob engine to SGLang-D that supports three denoising strategies and is compatible with Sequence Parallelism (SP).

Design Principles

• Stateless scheduler: All intermediate rollout state lives in a RolloutSessionData dataclass attached to the batch object (batch._rollout_session_data). The scheduler itself stores nothing, making it safe for concurrent use.
• Composable mixins: Rollout logic is encapsulated in SchedulerRLMixin (core) and SchedulerRLDebugMixin (debug tensor collection), mixed into FlowMatchEulerDiscreteScheduler alongside the existing SchedulerMixin and ConfigMixin.
• Opt-in activation: Rollout is triggered by SamplingParams.rollout = True. When disabled, no additional code paths execute and the scheduler step is unchanged.

Module Structure

runtime/post_training/
  rl_dataclasses.py            # RolloutSessionData, RolloutDebugTensors, RolloutTrajectoryData
  scheduler_rl_mixin.py        # SchedulerRLMixin  (core log-prob engine)
  scheduler_rl_debug_mixin.py  # SchedulerRLDebugMixin  (debug tensor accumulation)

Data Flow

prepare_rollout(batch, pipeline_config)
  --> creates RolloutSessionData on batch

for each denoising step:
  scheduler.step(batch=batch, ...)
    --> flow_sde_sampling(batch, model_output, sample, sigma, sigma_next, generator)
        --> computes prev_sample, log_prob_local_sum, local_elem_count
    --> append_local_rollout_log_probs(batch, ...)

after denoising loop:
  collect_rollout_log_probs(batch)      --> batch.rollout_trajectory_data.rollout_log_probs
  collect_rollout_debug_tensors(batch)  --> batch.rollout_trajectory_data.rollout_debug_tensors  (optional)
  release_rollout_resources(batch)      --> batch._rollout_session_data = None

2. Features, API & Reliability

Supported Rollout Strategies

API Parameters (SamplingParams)

Output

When rollout=True, the batch's rollout_trajectory_data (type RolloutTrajectoryData) is populated:

• rollout_log_probs: Tensor [B, T] -- per-step log-probabilities, reduced as global_sum / global_count across SP ranks.
• rollout_debug_tensors (when rollout_debug_mode=True): RolloutDebugTensors containing per-step [B, T, ...] tensors for variance_noise, prev_sample_mean, noise_std_dev, and model_output.

These are available through the OpenAI-compatible image/video API via the existing extra-fields mechanism.

Reliability Testing

Three levels of testing validate correctness:

1. Unit tests (test_scheduler_rollout_ode_unit.py):

   • ODE step determinism: confirms no variance noise is sampled.
   • Debug tensor shapes: verifies noise_std_dev is zero and shapes are consistent.
   • SDE/CPS FlowGRPO alignment: each single step is compared against a standalone FlowGRPO reference implementation across 4 seeds. Max absolute difference <= 1e-6 for prev_sample, prev_sample_mean, and noise_std_dev.

2. ODE-vs-non-rollout end-to-end (Z-Image-Turbo, 1024x1024, seed=42): ODE rollout produces bit-identical output images to the standard non-rollout denoising path across all parallelism configs (TP1/SP2, TP2/SP1, TP1/SP1+CFGP).

3. SDE/CPS parallel consistency (Z-Image-Turbo, 1024x1024, seed=42): Intermediate tensors (variance noise, noise std, prev sample mean, model output) are compared across parallelism configs against a single-GPU reference, with all steps and first-step metrics reported.

3. Detailed Test Results

3.1 ODE vs Non-Rollout (Bit-Exact Alignment)

Model: Tongyi-MAI/Z-Image-Turbo, seed=42, 1024x1024, 2 GPUs.

All configs produce pixel-identical images, confirming the ODE path introduces zero numerical divergence from the standard denoising pipeline.

3.2 SDE/CPS Parallel Consistency

Model: Tongyi-MAI/Z-Image-Turbo, seed=42, noise_level=0.5, 9 steps, 2 GPUs.
Reference: single-GPU (TP1, SP1, no CFGP parallel).

SDE mode -- key metrics (all steps):

CPS mode -- key metrics (all steps):

Key observations across both modes:

• variance_noise and noise_std_dev are bit-exact across all configs -- the SP-aware noise generation correctly reproduces the single-GPU noise stream.
• prev_sample_mean and model_output show small accumulated differences in TP2/SP1 and CFGP-parallel configs, caused by non-deterministic floating-point reduction order in the transformer (not the rollout engine). First-step differences are negligible (cosine > 0.9999).
• SP2 (sequence parallel on the latent axis) is bit-exact with the single-GPU reference on all tensors.

4. Summary

• The rollout log-prob engine correctly implements SDE, CPS, and ODE strategies as validated by unit tests aligned with FlowGRPO reference formulas.
• ODE mode produces bit-identical outputs to the standard non-rollout path across all parallelism configurations, confirming zero regression on existing functionality.
• Noise generation is bit-exact under SP -- the shard_latents_for_sp-based noise generation guarantees identical random streams regardless of parallelism.
• Small prev_sample_mean divergences under TP2/CFGP-parallel are expected floating-point non-determinism from the DiT forward pass, not the rollout engine. These accumulate over steps but remain within cosine similarity > 0.996.
• All rollout state is batch-scoped via RolloutSessionData -- the scheduler is fully stateless and safe for multi-request serving.