Repository navigation
[Refactor] diffusion: Rollout controller as a scheduler-step plugin instead of mixins - #44603
Merged
Merged
Conversation
…troller.py Pure file rename (git R100); the proof tool has no file-rename primitive. Refactor-Kind: non_mechanical_provable Signed-off-by: rockdu <kangrdu@gmail.com>
…er.py Refactor-Kind: mechanical_provable Signed-off-by: rockdu <kangrdu@gmail.com>
Refactor-Kind: mechanical_provable Signed-off-by: rockdu <kangrdu@gmail.com>
Refactor-Kind: mechanical_provable Signed-off-by: rockdu <kangrdu@gmail.com>
… RolloutSessionData Refactor-Kind: non_mechanical_provable Signed-off-by: rockdu <kangrdu@gmail.com>
Refactor-Kind: mechanical_provable Signed-off-by: rockdu <kangrdu@gmail.com>
Rockdu
force-pushed
the
rockdu/diffusion-rollout-plugin
branch
from
October 15, 2026 23:49
5a3eb3b to
5d8d7c0
Compare
…lloutController and drop SchedulerRLMixin Refactor-Kind: non_mechanical_provable Signed-off-by: rockdu <kangrdu@gmail.com>
…stead of branching on batch Refactor-Kind: non_mechanical_provable Co-authored-by: WenhaoZhang <42087078+niehen6174@users.noreply.github.com> Signed-off-by: rockdu <kangrdu@gmail.com>
… of the denoise loops Refactor-Kind: non_mechanical_provable Co-authored-by: WenhaoZhang <42087078+niehen6174@users.noreply.github.com> Signed-off-by: rockdu <kangrdu@gmail.com>
… hook functions Refactor-Kind: non_mechanical_provable Co-authored-by: WenhaoZhang <42087078+niehen6174@users.noreply.github.com> Signed-off-by: rockdu <kangrdu@gmail.com>
…ontroller Refactor-Kind: non_mechanical_provable Co-authored-by: WenhaoZhang <42087078+niehen6174@users.noreply.github.com> Signed-off-by: rockdu <kangrdu@gmail.com>
Rockdu
force-pushed
the
rockdu/diffusion-rollout-plugin
branch
from
October 15, 2026 23:58
5d8d7c0 to
3598b8c
Compare
3 of 6 tasks
… miles-diffusion Refactor-Kind: non_mechanical_provable Signed-off-by: rockdu <kangrdu@gmail.com>
Rockdu
marked this pull request as ready for review
October 17, 2026 00:56
Rockdu
requested review from
AgainstEntropy,
BBuf,
HaiShaw,
kevin-mii,
mickqian,
niehen6174,
ping1jing2 and
yichiche
as code owners
October 17, 2026 00:56
This comment has been minimized.
This comment has been minimized.
Rockdu
force-pushed
the
rockdu/diffusion-rollout-plugin
branch
from
October 17, 2026 03:42
0bb53ce to
3567900
Compare
This comment has been minimized.
This comment has been minimized.
… keep trajectory timesteps on device Refactor-Kind: non_mechanical_provable Co-authored-by: Kangrui Du <89372739+Rockdu@users.noreply.github.com> Co-authored-by: JingwenGu0829 <75733630+JingwenGu0829@users.noreply.github.com> Signed-off-by: rockdu <kangrdu@gmail.com>
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
Rockdu
force-pushed
the
rockdu/diffusion-rollout-plugin
branch
from
October 17, 2026 07:25
9cf7390 to
e3354ce
Compare
This comment has been minimized.
This comment has been minimized.
Rockdu
changed the title
Rockdu
changed the title
Rockdu
changed the title
This comment has been minimized.
This comment has been minimized.
4 of 5 tasks
This comment has been minimized.
This comment has been minimized.
Rockdu
force-pushed
the
rockdu/diffusion-rollout-plugin
branch
from
October 17, 2026 09:35
c06bf48 to
30056b9
Compare
This comment has been minimized.
This comment has been minimized.
Rockdu
force-pushed
the
rockdu/diffusion-rollout-plugin
branch
2 times, most recently
from
October 17, 2026 09:45
c06bf48 to
5cd64bd
Compare
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
Collaborator
|
/tag-and-rerun-ci |
Rockdu
changed the title
# Conflicts: # python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py
2 of 4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Rollout stops being a scheduler / denoising-stage mixin: each latent stream of a rollout request gets a
FlowRolloutControllerwhosetransitionthe loop passes intoscheduler.step. Serving is untouched; rollout numerics do not change.It supersedes the design in #37940; @niehen6174 is co-author on the derived commits. @JingwenGu0829 found the issues fixed in commits 13 and 14 (Rockdu#4) and is co-author on commits 13 to 15.
How it works
FlowMatchEulerDiscreteScheduler.step(..., transition=None)runstransition or self._transition. Loops pass**rollout_step_kwargs(batch), which is empty for serving requests, so every scheduler keeps its serving call.rollout_hooks.py:prepare_rollout_if_enabled(before the SP shard),init_rollout_env_if_enabled,rollout_step_kwargs,finalize_rollout_if_enabled; all no-ops for serving.DenoisingStage, Cosmos3 and LTX-2. A loop without a scheduler (MiniMax H3, next PR) passes its ownbuild_controllerandode_transition.The RL model support in SGLang diffusion will then be like:
Review guide
Commits 2, 3, 4 and 6 are mechanical moves verified by the chain verifier (4/4 proofs) and commit 1 is a pure rename; skim them. Review commits 5, 7 to 11 and 13 to 15.
Tests
main(8x H200, SP 1 and 2, every returned tensor): Qwen-Image and Cosmos3 are bitwise identical; LTX-2.3 trajectory and log-probs are bitwise identical (its audio decode is nondeterministic onmaintoo). The only intended change: bf16 models now return fp32dit_trajectory.latentswith the same values. The A/B predates commits 13 to 15, which the e2e below covers.sglang-milesplus this chain: every suite passes, bitwise equal to its standard.Checklist
CI States
Latest PR Test (Base): ✅ Run #39336833020
Latest PR Test (Extra): ❌ Run #39336832991
Latest PR Test (AMD ROCm 10): ❌ Run #39336833006