[Rockdu]

Review guide

TL;DR: rollout stops being a scheduler mixin. It becomes a per-request controller (FlowRolloutController) whose transition is passed to scheduler.step(..., transition=...). The denoising-stage mixin becomes three hook functions. Rollout numerics do not change.

This is the sglang-miles port of the same 11-commit chain prepared for main. That main PR is not opened yet; this draft comes first for review. It supersedes the design in #37940 and credits @niehen6174 as co-author on the derived commits.

• Commits 2, 3, 4 and 6 are mechanical_provable (refactor(move)): each reproduces byte-for-byte from its parent with the mechanical-refactor-verify primitives, and the chain verifier passes (4/4 proofs). Skim them.
• Commit 1 is a pure file rename (git R100); the proof tool has no file-rename primitive, so it is labeled non_mechanical_provable.
• Review commits 5 and 7 to 11.

Motivation

Rollout used to live inside the scheduler and the denoising stage as mixins (SchedulerRLMixin, SchedulerRLDebugMixin, RolloutDenoisingMixin). step() branched on batch.rollout, and every denoise loop recorded x_t itself. That shape cannot be plugged into a different loop or a second latent stream, such as the audio stream that comes next. After this PR, a rollout is a self-contained object that any loop can hand to any FlowMatch Euler step().

Modifications

#	Commit	Refactor-Kind	Proof
1	rename scheduler_rl_mixin.py -> flow_rollout_controller.py	non_mechanical_provable	git R100
2	move SchedulerRLDebugMixin into flow_rollout_controller.py	mechanical_provable	PASS
3	extract the FlowMatch Euler update into _transition	mechanical_provable	PASS
4	move RolloutSessionData into flow_rollout_controller.py	mechanical_provable	PASS
5	prepare the session methods to move onto RolloutSessionData (explicit step_index, retype self)	non_mechanical_provable	review
6	move the session methods into RolloutSessionData	mechanical_provable	PASS
7	RolloutSessionData -> per-request FlowRolloutController on Req.rollout_controllers["visual"], holding the loop's sigma grid; drop SchedulerRLMixin	non_mechanical_provable	review
8	pass transition=rollout_transition(batch) to scheduler.step instead of branching on batch	non_mechanical_provable	review
9	record x_t inside the rollout transition instead of in each denoise loop; finalize reads the trajectory and sigmas from the controller	non_mechanical_provable	review
10	RolloutDenoisingMixin -> prepare_rollout_if_enabled(build_controller=...) / init_rollout_env_if_enabled / finalize_rollout_if_enabled; the default builder constructs the controller from batch.scheduler	non_mechanical_provable	review
11	optional ode_transition: a model with its own update (MiniMax H3, next PR) passes it for ODE steps; None keeps the Euler update	non_mechanical_provable	review

How it works now:

• FlowMatchEulerDiscreteScheduler.step(..., transition=None) calls transition or self._transition. The scheduler instance is never patched, so stage-shared schedulers (Krea-2, Qwen-Image-Layered) stay safe. The existing fp32 upcast at the top of step() is untouched.
• FlowRolloutController is created once per request at the old prepare site, before the SP shard. It owns the noise buffer, log-probs, debug tensors, the trajectory, and its SP shard/gather closures.
• The controller needs only the loop's sigma grid, not a scheduler object. prepare_rollout_if_enabled builds it with _build_scheduler_rollout_controller for loops that step a FlowMatchEulerDiscreteScheduler; a loop with its own update and no scheduler (MiniMax H3, next PR) passes its own build_controller(batch=, pipeline_config=), which also sets its ode_transition. No loop in this PR passes one.
• Generic DenoisingStage, Cosmos3 and LTX-2 pass the transition at their step() call. LTX-2 drops its _step_index / _scheduler_step_kwargs workaround. MiniMax H3 keeps its private rollout; moving it onto this plugin is the follow-up PR.

sglang-miles-specific adaptations (everything else is identical to the main chain):

• Kept sglang-miles's trajectory semantics: timesteps cover every step, latents follow rollout_return_step_indices, and latent_step_indices is recorded. They now live in FlowRolloutController._maybe_append_dit_trajectory_step and finalize_rollout_if_enabled.
• sglang-miles has no SigmaScheduleMixin, randn_tensor or _run_denoising_loop. The resolutions follow the base.
• test_timestep_filters_gate_sde_and_trajectory asserted filtered timesteps, which does not match sglang-miles's own latent_step_indices change. The base's circular import (below) hid this. The test is updated to the actual semantics.

Accuracy Tests

This PR does not change outputs. One difference is intended: for models whose latents are bf16, the returned dit_trajectory.latents are now fp32. The values are unchanged, because x_t is recorded after step()'s own upcast. miles_diffusion casts trajectory latents to fp32 before use.

GPU A/B of this branch (tip) against sglang-miles on 8x H200: same seed, debug mode on, every returned tensor compared, dit_trajectory.sigmas included.

Model	SP	Cases	Result
Qwen-Image	1, 2	sde, cps, ode, filtered sde trajectory	bitwise identical
Cosmos3	1, 2	sde, ode	bitwise identical
LTX-2.3	1, 2	sde, cps, ode	bitwise identical except the audio waveform. The waveform also differs between two identical runs of the same tree (nondeterministic audio decode); trajectory and log-probs are identical.

miles_diffusion e2e on this sglang-miles port (tip 8d835f684c): every suite passed with no standard re-recorded (radixark/miles_diffusion#313, full run with --match-all-labels --continue-on-error).

Suite	Tests	Result
stage-c 3-GPU	SD3 NFT, MiniMax H3 t2va GRPO, MiniMax H3 t2va SFT, SD3 OCR GRPO	4/4 pass
stage-c 5-GPU	Wan2.2 GRPO, Cosmos3 GRPO, LTX-2.3 GRPO, Qwen-Image GRPO, Krea-2 NFT async, SP attention parity	6/6 pass
stage-a/b CPU and GPU	unit and fast-GPU tests	all pass

Unit tests

Run on test_scheduler_rollout_unit.py, test_cosmos3_rollout.py and test_rollout_api.py:

• On sglang-miles, two of these files fail at import with a pre-existing circular import (scheduler_rl_mixin <-> schedule_batch).
• Every commit from 7 onward passes: 52 passed and 1 skipped at the tip.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations. Not needed: no user-facing change.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed. The A/B and the e2e above cover accuracy; no speed impact expected.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #38683323869
Latest PR Test (Extra): ❌ Run #38683323858
Latest PR Test (AMD ROCm 10): ❌ Run #38683323867