[diffusion] Support RL rollout for the Wan pipeline via a per-request scheduler switch

[zhihengy]

Motivation

main already ships the diffusion RL rollout infrastructure (POST /rollout/generate, SchedulerRLMixin, RolloutDenoisingMixin), but Wan is not wired into it: WanPipeline uses a bare FlowUniPCMultistepScheduler, so rollout requests cannot run SDE sampling with log-probs. Serving and rollout need different schedulers on one engine — Wan serves with UniPC while the rollout SDE/log-prob path requires first-order flow-match Euler — and the choice must be per request, since one engine serves both rollout (rollout=True) and eval (rollout=False) requests within a training run.

Modifications

Same dormant pattern as RolloutDenoisingMixin on the standard DenoisingStage: rollout logic lives under post_training/ and only activates when a request sets rollout=True; a pure-serving engine assembles and runs exactly as on main.

• post_training/rollout_timestep_mixin.py (new) — RolloutTimestepPreparationMixin: per-request scheduler selection on batch.rollout, plus a sanity check of the flow-match Euler convention timesteps == sigmas[:-1] * num_train_timesteps that the SDE/log-prob math relies on (fails loudly on a mismatched scheduler; logs the active rollout scheduler once per process).
• pipelines_core/stages/timestep_preparation.py — TimestepPreparationStage mixes in the above and takes an optional rollout_scheduler (default None; existing pipelines unchanged). Downstream stages read batch.scheduler, so this stage is the single binding point.
• pipelines/wan_pipeline.py — binding only: the scheduler module stays a bare UniPC (identical to main); the pipeline passes a FlowMatchEulerDiscreteScheduler (same launch-time flow_shift) as the stage's rollout_scheduler.

Other pipelines whose serving scheduler is not first-order Euler (Wan I2V, SANA-class) can adopt the same one-line binding.

Accuracy Tests

• Verified locally: rollout=False binds UniPC, rollout=True binds flow-match Euler; pipelines without a rollout_scheduler are unaffected; the convention check rejects a mismatched scheduler. Full multimodal_gen/test/unit suite unchanged before/after.
• Exercised downstream: Wan2.2-T2V-A14B GRPO training (SDE log-probs, per-step rollout filters, dual-expert weight sync) runs on this wiring, with rollout-vs-trainer log_prob agreement at the bf16 noise floor (~1e-5).

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.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #28832550262
Latest PR Test (Extra): ❌ Run #28832550138

[gemini-code-assist]

Code Review

This pull request introduces the WanRolloutScheduler and WanTimestepPreparationStage to support both normal inference (using UniPC) and RL rollout (using Euler SDE) in the Wan video diffusion pipeline. The review feedback highlights critical improvements for concurrent serving environments: first, to prevent race conditions, the shared template scheduler should not be mutated directly, and a request-specific scheduler instance should be used instead; second, to avoid log flooding, the logging guard should be changed from an instance variable to a class-level variable since schedulers are cloned per request.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Mutating the shared template scheduler (self.scheduler) in WanTimestepPreparationStage.forward can lead to race conditions in concurrent serving environments where multiple requests are processed simultaneously.

Instead, you should first obtain the request-specific scheduler instance using get_or_create_request_scheduler and then call prepare_for_batch on that request-specific instance. This ensures that the state is isolated per request.

Suggested change

	def forward(self, batch, server_args):
	self.scheduler.prepare_for_batch(batch)
	return super().forward(batch, server_args)
	def forward(self, batch, server_args):
	from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import (
	get_or_create_request_scheduler,
	)
	scheduler = get_or_create_request_scheduler(batch, self.scheduler)
	scheduler.prepare_for_batch(batch)
	return super().forward(batch, server_args)

[gemini-code-assist]

Since the scheduler is copied/cloned per request (via get_or_create_request_scheduler), instance variables like self._logged_rollout_euler_check will be reset to False for each new request's scheduler instance. This will cause the info log to be printed once per request, potentially flooding the logs in high-throughput environments.

Using a class-level variable instead of an instance variable ensures that the log is printed only once globally.

Suggested change

	class WanRolloutScheduler(SchedulerRLMixin):
	"""Use UniPC for normal Wan inference and Euler SDE for RL rollout."""
	def __init__(self, shift: float | None):
	self.unipc_scheduler = FlowUniPCMultistepScheduler(shift=shift)
	self.euler_scheduler = FlowMatchEulerDiscreteScheduler(
	shift=1.0 if shift is None else shift
	)
	self._active_scheduler = self.unipc_scheduler
	self._logged_rollout_euler_check = False
	class WanRolloutScheduler(SchedulerRLMixin):
	"""Use UniPC for normal Wan inference and Euler SDE for RL rollout."""
	_logged_rollout_euler_check = False
	def __init__(self, shift: float | None):
	self.unipc_scheduler = FlowUniPCMultistepScheduler(shift=shift)
	self.euler_scheduler = FlowMatchEulerDiscreteScheduler(
	shift=1.0 if shift is None else shift
	)
	self._active_scheduler = self.unipc_scheduler

[gemini-code-assist]

Update the check to use the class-level _logged_rollout_euler_check variable to ensure the log is printed only once globally.

Suggested change

	if not self._logged_rollout_euler_check:
	logger.info(
	"Wan rollout using FlowMatchEulerDiscreteScheduler "
	"(timesteps dtype=%s, sigmas dtype=%s, max_abs_diff=%.6g)",
	timesteps.dtype,
	sigmas.dtype,
	max_abs_diff,
	)
	self._logged_rollout_euler_check = True
	if not WanRolloutScheduler._logged_rollout_euler_check:
	logger.info(
	"Wan rollout using FlowMatchEulerDiscreteScheduler "
	"(timesteps dtype=%s, sigmas dtype=%s, max_abs_diff=%.6g)",
	timesteps.dtype,
	sigmas.dtype,
	max_abs_diff,
	)
	WanRolloutScheduler._logged_rollout_euler_check = True

[zhihengy]

/tag-run-ci-label

[mickqian]

could we make sure only one scheduler is initialized and passed at runtime?