[Diffusion] Return scheduler sigmas snapshot in rollout dit_trajectory

[Rockdu]

What

• Add sigmas to RolloutDitTrajectory: a [T+1] snapshot of scheduler.sigmas (post-shift, includes the terminal 0), captured when the trajectory is finalized and returned alongside timesteps in POST /rollout/generate responses.

Why

Training-side consumers currently recompute sigmas from trajectory timesteps as timesteps / num_train_timesteps. That round-trips sigma * 1000 / 1000 and drifts ULPs, which amplifies into observable log-prob differences when the division happens in bf16. The scheduler already holds the exact per-request schedule (set_timesteps rebuilds self.sigmas for the request's actual num_inference_steps), so return it directly.

Semantics worth noting for reviewers:

• sigmas is schedule metadata, not per-entry data: it stays the full [T+1] schedule and is not affected by rollout_return_step_indices filtering. When filtering is used, consumers must align by absolute step index rather than zipping positionally with the filtered timesteps.
• It is batch-global: with num_outputs_per_prompt > 1 all samples of a request share one scheduler schedule, so the per-sample slicer passes it through unsliced (same pattern as timesteps). This also avoids _extract_single_sample_tensor mis-slicing a 1-D schedule whose length happens to equal batch_size.
• Samplers that call set_timesteps more than once per generation (multi-stage pipelines) would snapshot the last schedule; this matches the existing timesteps behavior and single-stage rollout is the only supported RL path today.

Validation

• python -m pytest python/sglang/multimodal_gen/test/unit/test_rollout_api.py -q: 34 passed (H200).
• E2E on H200: launched stabilityai/stable-diffusion-3.5-medium, 10-step rollout via POST /rollout/generate with rollout_return_dit_trajectory=true; response dit_trajectory.sigmas is fp32 (11,) with terminal 0, and sigma * 1000 bitwise-equals the returned timesteps.

Files

• runtime/post_training/rl_dataclasses.py — add the sigmas field.
• runtime/post_training/rollout_denoising_mixin.py — snapshot batch.scheduler.sigmas at trajectory finalize.
• runtime/entrypoints/post_training/rollout_api.py — pass sigmas through per-sample extraction and response serialization.
• test/unit/test_rollout_api.py — extend serialization and batched-response tests.

Checklist

• pre-commit run --files <changed files> passes (all hooks, no auto-fixes needed)
• Added/updated tests for new behaviour
• pytest -x is green (scoped to test_rollout_api.py, 34 passed)
• No launch flags changed
• No public flag added
• No example added

---

CI States

Latest PR Test (Base): ✅ Run #30410262021
Latest PR Test (Extra): ❌ Run #30410261848

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[AgainstEntropy]

/tag-and-rerun-ci

[Rockdu]

TODO: ablation and switch to scheduler.timesteps as well