[Diffusion] Refactor CFG parallel for extensibility and performance

[TKONIY]

Summary

Refactor CFGParallelMixin to improve performance, reduce code complexity, and enable extensibility for multi-output models (e.g., world models producing both video and action predictions).

Problem

Current CFGParallelMixin has three issues:

1. Redundant communication: Every denoising step does all_gather + broadcast. After all_gather, both ranks already have identical positive and negative predictions — the broadcast in scheduler_step_maybe_with_cfg is unnecessary.

2. Code complexity: predict_noise_maybe_with_cfg returns valid result only on rank 0, None on rank 1. scheduler_step_maybe_with_cfg has rank-0-only step + broadcast. All downstream code must handle None/rank-branching.

3. Limited extensibility: predict_noise_maybe_with_cfg assumes single tensor output. Models with multi-output (e.g., DreamZero world model returning (video_pred, action_pred)) cannot use the mixin without overriding the entire method.

Changes

Change 1: Remove redundant broadcast

• After all_gather, all ranks compute CFG combine + scheduler step locally
• Eliminates one broadcast per denoising step
• predict_noise_maybe_with_cfg returns valid tensor on all ranks (not just rank 0)
• scheduler_step_maybe_with_cfg simplified: just calls scheduler_step on all ranks

Change 2: Tuple output support via combine_cfg_noise override

• predict_noise() can now return tuple[Tensor, ...] for multi-output models
• Mixin internally normalizes to tuple via _wrap/_unwrap helpers, processes uniformly
• combine_cfg_noise() accepts tuple, applies CFG formula per-element by default
• Multi-output models override combine_cfg_noise() for custom per-element logic:

class MyWorldModelPipeline(nn.Module, CFGParallelMixin):
    def predict_noise(self, **kwargs):
        video_pred, action_pred = self.transformer(**kwargs)
        return (video_pred, action_pred)

    def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, scale, normalize):
        (video_pos, action_pos) = positive_noise_pred
        (video_neg, action_neg) = negative_noise_pred
        (video_combined,) = super().combine_cfg_noise((video_pos,), (video_neg,), scale, normalize)
        return (video_combined, action_pos)  # action: positive only, no CFG

Correctness Argument

After all_gather, both ranks have bit-identical copies of positive/negative predictions. CFG combine and scheduler step are deterministic (pure arithmetic, no randomness). Same inputs + same ops = same results on both ranks.

Benchmark

Setup:

• Model: Qwen/Qwen-Image-2512 (25B)
• Hardware: NVIDIA RTX PRO 6000 Blackwell (98GB VRAM), PCIe Gen5 x16, PIX topology
• Parameters: cfg_scale=4.0, steps=20, 1024x1024, seed=142
• Main branch: cb6d012d (latest origin/main)

Performance

Test	Branch	CFG Parallel	GPUs	Total Time (ms)	vs main
1	main	Off	1	9823	—
2	main	On (size=2)	2	5284	—
3	refactor	Off	1	9797	-0.3%
4	refactor	On (size=2)	2	5205	-1.5%

No performance regression. Broadcast cost is negligible on this hardware/model (PCIe Gen5, ~0.1ms vs ~250ms transformer forward per step). The value of this PR is code simplification and extensibility, not raw speed.

Correctness

Comparison	MD5 Match
baseline_no_cfg vs refactor_no_cfg	Identical (972a8c57)
baseline_cfg_parallel vs refactor_cfg_parallel	Identical (c218a5aa)

Bit-identical outputs confirm the refactor produces exactly the same results.

Test plan

• Benchmark: Qwen-Image-2512 (25B) on 1 GPU and 2 GPU CFG parallel — no regression
• Correctness: bit-identical images (MD5 match) between main and refactor
• Unit test: tests/diffusion/distributed/test_cfg_parallel.py now verifies all CFG ranks receive the same combined output and match the sequential CFG baseline
• Run existing CI tests

Next step

• After this PR is merged, refactor LTX2 and LTX2 I2V (video + audio) with the refactor CFG parallel framework.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7f3b05931d

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[wtomin]

I agree with philosophy of this refactoring PR, as it brings better extensibility.

CFG combine and scheduler step are deterministic (pure arithmetic, no randomness).

I am afraid some diffusion schedulers are not deterministic, for example, DDPM solver. Although in vLLM-Omni, most of diffusion models use deterministic schedulers, like flow matching.

Therefore, I suggest to better handle the randomness of scheduler_step_maybe_with_cfg

[wtomin]

Bagel, GLM-Image, dreamid-omni, nextstep1.1, and LTX-2 do not inherit from CFGParallelMixmin, but implement CFG Parallel from scratch. Please check if your PR affect them.

[TKONIY]

Bagel, GLM-Image, dreamid-omni, nextstep1.1, and LTX-2 do not inherit from CFGParallelMixmin, but implement CFG Parallel from scratch. Please check if your PR affect them.

Thanks for the thorough review! Here's what's been addressed:

1. Non-deterministic scheduler (generator support)

Done in a5f2de3d — scheduler_step() and scheduler_step_maybe_with_cfg() now accept an optional generator param, passed through to the scheduler's step() call.

2 & 3. Doc updates (d53df5a2)

• Added generator=None to the composite scheduler example (VideoAudioScheduler.step()) and the diffuse() call site
• Added a note about explicitly passing generator=torch.Generator(device).manual_seed(seed) for non-deterministic schedulers like DDPM — adopted your suggested wording

4. Impact on Bagel / GLM-Image / DreamID-Omni / NextStep 1.1 / LTX-2

Checked all five:

Model	Inherits Mixin?	Affected?
Bagel	No	No — fully independent CFG impl
GLM-Image	No	No — uses parallel_state directly
NextStep 1.1	No	No — uses parallel_state directly
DreamID-Omni	Yes	No — only calls combine_cfg_noise() with plain tensors; _wrap/_unwrap preserves backward compat
LTX-2	Yes	No — same as above; uses its own predict_noise_av_maybe_with_cfg

DreamID-Omni and LTX-2 inherit CFGParallelMixin but only use combine_cfg_noise(), which now handles both plain Tensor and tuple[Tensor, ...] transparently. Verified with standalone unit tests covering both input types.

btw, I have an on-going draft PR refactoring LTX-2 on the new cfg_parallel.

• [Diffusion] Refactor LTX2 to use unified CFG parallel framework and enable per_request_scheduler #2078
• [Diffusion] Refactor LTX2 to use unified CFG parallel framework #2160

[princepride]

PTAL.

[nussejzz]

Bagel, GLM-Image, dreamid-omni, nextstep1.1, and LTX-2 do not inherit from CFGParallelMixmin, but implement CFG Parallel from scratch. Please check if your PR affect them.

@wtomin This shouldn't affect bagel, because that's how I implemented it. #1695

[wtomin]

LGTM

[princepride]

LGTM