[Refactor] Extend CFG Parallel to support 3 or 4 branch dispatch across M GPUs (#2423)

docs/design/feature/cfg_parallel.md

@@ -25,7 +25,9 @@ In standard Classifier-Free Guidance, each diffusion step requires two forward p
 1. **Positive/Conditional**: Guided by the text prompt
 2. **Negative/Unconditional**: Typically using empty or negative prompt
 
-CFG-Parallel eliminates this bottleneck by distributing the two forward passes across different GPU ranks, allowing them to execute simultaneously rather than sequentially.
+Some models require 3 or more CFG branches (see [N-Branch CFG](#n-branch-cfg-3-branches)).
+
+CFG-Parallel eliminates this bottleneck by distributing the forward passes across different GPU ranks, allowing them to execute simultaneously rather than sequentially.
 
 ### Architecture
 
@@ -33,9 +35,11 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
 
 | Method | Purpose | Automatic Behavior |
 |--------|---------|-------------------|
-| [`predict_noise_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Predict noise with CFG | Detects parallel mode, distributes computation, gathers results |
+| [`predict_noise_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Predict noise with 2-branch CFG | Detects parallel mode, distributes computation, gathers results |
+| [`predict_noise_with_multi_branch_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Predict noise with N-branch CFG | Round-robin dispatches N branches across M GPUs |
 | [`scheduler_step_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Step scheduler | All ranks step locally (no broadcast needed) |
-| [`combine_cfg_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Combine positive/negative | Applies CFG formula with optional normalization |
+| [`combine_cfg_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Combine 2-branch predictions | Applies CFG formula with optional normalization |
+| [`combine_multi_branch_cfg_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Combine N-branch predictions | Override for custom multi-branch combine logic |
 | [`predict_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Forward pass wrapper | Override for custom transformer calls |
 | [`cfg_normalize_function()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Normalize CFG output | Override for custom normalization |
 
@@ -57,6 +61,22 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
 
 - All ranks compute the scheduler step locally — no broadcast needed because `predict_noise_maybe_with_cfg` already ensures all ranks have identical noise predictions after `all_gather` + local combine.
 
+### N-Branch CFG (3+ branches)
+
+Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches, DreamID Omni uses 4 branches.
+
+`predict_noise_with_multi_branch_cfg()` handles these by automatically dispatching N branches across M GPUs using round-robin (rule: branch `i` → rank `i % M`):
+
+| Branches (N) | GPUs (M) | Dispatch |
+|:---:|:---:|:---|
+| 3 | 2 | `[[0, 2], [1]]` |
+| 3 | 3 | `[[0], [1], [2]]` |
+| 4 | 2 | `[[0, 2], [1, 3]]` |
+| 4 | 3 | `[[0, 3], [1], [2]]` |
+| 4 | 4 | `[[0], [1], [2], [3]]` |
+
+When a rank handles multiple branches, it runs them sequentially. After `all_gather`, all ranks execute `combine_multi_branch_cfg_noise()` locally, producing identical results.
+
 ---
 
 ## Step-by-Step Implementation
@@ -98,6 +118,7 @@ class YourModelPipeline(nn.Module, CFGParallelMixin):
 - `positive_kwargs`: transformer arguments for conditional (text-guided) prediction
 - `negative_kwargs`: transformer arguments for unconditional prediction (set to `None` if CFG disabled)
 - For image editing pipelines, add `output_slice=image_seq_len` to extract the generative image portion
+- For models with 3+ CFG branches, see [Multi-Branch CFG](#multi-branch-cfg-3-branches) in the Customization section
 
 ### Step 2: Call `diffuse`
 
@@ -171,20 +192,42 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin):
 ```
 
 
-### Override `combine_cfg_noise()` for Multi-Output Models
+### Multi-Branch CFG (3+ branches)
 
-When `predict_noise()` returns a tuple (e.g., video + audio), the default `combine_cfg_noise()` applies CFG to every element. Override it to apply different logic per element — for example, CFG on video but positive-only on audio:
+For models with 3 or more CFG branches, use `predict_noise_with_multi_branch_cfg()` instead of `predict_noise_maybe_with_cfg()`, and override `combine_multi_branch_cfg_noise()` for custom combine logic. This interface also works for standard 2-branch CFG — just pass 2 branches in `branches_kwargs`.
+
+**Example (3-branch with dual guidance scale):**
 
 ```python
-class MyVideoAudioPipeline(nn.Module, CFGParallelMixin):
-    def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, scale, normalize):
-        (video_pos, audio_pos) = positive_noise_pred
-        (video_neg, audio_neg) = negative_noise_pred
-        video_combined = super().combine_cfg_noise(video_pos, video_neg, scale, normalize)
-        return (video_combined, audio_pos)  # audio: positive only, no CFG
+class YourMultiBranchPipeline(nn.Module, CFGParallelMixin):
+    def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize=False):
+        text_scale = true_cfg_scale["text"]
+        image_scale = true_cfg_scale["image"]
+        pos, ref, uncond = predictions
+        return uncond + image_scale * (ref - uncond) + text_scale * (pos - ref)
+
+    def diffuse(self, ...):
+        for i, t in enumerate(timesteps):
+            positive_kwargs = {...}   # conditional prompt
+            ref_neg_kwargs = {...}    # negative prompt + reference
+            uncond_kwargs = {...}     # unconditional
+
+            noise_pred = self.predict_noise_with_multi_branch_cfg(
+                do_true_cfg=do_true_cfg,
+                true_cfg_scale={"text": text_guidance_scale, "image": image_guidance_scale},
+                branches_kwargs=[positive_kwargs, ref_neg_kwargs, uncond_kwargs],
+            )
+            latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
+
+        return latents
 ```
 
-This also requires `predict_noise()` to return a tuple (see [Override predict_noise](#override-predict_noise-for-custom-transformer-calls) above).
+### Override Combine Functions
+
+There are two combine functions for different scenarios:
+
+- **`combine_cfg_noise()`** — Used by `predict_noise_maybe_with_cfg()`. Override when `predict_noise()` returns a tuple (e.g., video + audio) and you need per-element CFG logic.
+- **`combine_multi_branch_cfg_noise()`** — Used by `predict_noise_with_multi_branch_cfg()`. Override to implement custom multi-branch combine formulas (see [Multi-Branch CFG](#multi-branch-cfg-3-branches) above).
 
 ### Implement a Composite Scheduler for Multi-Output Models
 
@@ -303,4 +346,5 @@ Adding CFG-Parallel support:
 
 1. ✅ **Create mixin** - Inherit from `CFGParallelMixin` and implement `diffuse()` method
 2. ✅ **(Optional) Customize** - Override `predict_noise()` or `cfg_normalize_function()` for custom behavior
-3. ✅ **Test** - Verify with `--cfg-parallel-size 2` and compare performance
+3. ✅ **(Optional) Multi-branch** - For 3+ branch models, use `predict_noise_with_multi_branch_cfg()` and override `combine_multi_branch_cfg_noise()`
+4. ✅ **Test** - Verify with `--cfg-parallel-size 2` (or 3/4 for multi-branch) and compare performance

docs/user_guide/diffusion_features.md

@@ -118,7 +118,7 @@ The following tables show which models support each feature:
 | **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
 | **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **OmniGen2** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
 | **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
 | **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |

examples/offline_inference/image_to_image/image_edit.py

@@ -297,8 +297,8 @@ def parse_args() -> argparse.Namespace:
         "--cfg-parallel-size",
         type=int,
         default=1,
-        choices=[1, 2],
-        help="Number of GPUs used for classifier free guidance parallel size.",
+        choices=[1, 2, 3],
+        help="Number of GPUs used for classifier free guidance parallel size (max 3 branches).",
     )
     parser.add_argument(
         "--enforce-eager",

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -36,8 +36,8 @@ def parse_args() -> argparse.Namespace:
         "--cfg-parallel-size",
         type=int,
         default=1,
-        choices=[1, 2],
-        help="Number of GPUs used for classifier free guidance parallel size.",
+        choices=[1, 2, 3, 4],
+        help="Number of GPUs used for classifier free guidance parallel size (max 4 branches).",
     )
     parser.add_argument(
         "--video-negative-prompt",

tests/diffusion/distributed/test_cfg_parallel.py

@@ -2,8 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Unit tests for CFG (Classifier-Free Guidance) parallel functionality.
 
-This test verifies that predict_noise_maybe_with_cfg produces numerically
-equivalent results with and without CFG parallel using fixed random inputs.
+This test verifies that predict_noise_maybe_with_cfg and
+predict_noise_with_multi_branch_cfg produce numerically equivalent results
+with and without CFG parallel using fixed random inputs.
 """
 
 import os
@@ -429,3 +430,340 @@ def test_predict_noise_without_cfg(dtype: torch.dtype):
     assert noise_pred.shape == (1, 4, 16, 16)
 
     print(f"✓ Test passed: predict_noise without CFG (dtype={dtype})")
+
+
+class MultiBranchTestPipeline(CFGParallelMixin):
+    """Test pipeline with custom 3-branch combine logic (like OmniGen2)."""
+
+    def __init__(self, in_channels: int = 4, hidden_dim: int = 128, seed: int = 42):
+        torch.manual_seed(seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(seed)
+
+        self.transformer = SimpleTransformer(in_channels, hidden_dim)
+
+        torch.manual_seed(seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(seed)
+        for param in self.transformer.parameters():
+            torch.nn.init.normal_(param, mean=0.0, std=0.02)
+
+    def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize=False):
+        """N-branch combine with weighted sum for testing.
+
+        - 2-branch: standard CFG formula (true_cfg_scale is float)
+        - 3-branch: OmniGen2-style dual guidance scale (true_cfg_scale is dict)
+        - 4-branch: DreamID-style weighted sum (true_cfg_scale is dict)
+        """
+        if len(predictions) == 4:
+            text_scale = true_cfg_scale["text"]
+            image_scale = true_cfg_scale["image"]
+            vid_ref_scale = true_cfg_scale["vid_ref"]
+            pos, neg, vid_neg, audio_neg = predictions
+            combined = (
+                audio_neg
+                + vid_ref_scale * (vid_neg - audio_neg)
+                + image_scale * (neg - vid_neg)
+                + text_scale * (pos - neg)
+            )
+        elif len(predictions) == 3:
+            text_scale = true_cfg_scale["text"]
+            image_scale = true_cfg_scale["image"]
+            pos, ref, uncond = predictions
+            combined = uncond + image_scale * (ref - uncond) + text_scale * (pos - ref)
+        else:
+            pos, neg = predictions[0], predictions[1]
+            combined = neg + true_cfg_scale * (pos - neg)
+
+        if cfg_normalize:
+            combined = self.cfg_normalize_function(pos, combined)
+        return combined
+
+
+def _test_multi_branch_parallel_worker(
+    local_rank: int,
+    world_size: int,
+    cfg_parallel_size: int,
+    dtype: torch.dtype,
+    test_config: dict,
+    result_queue: torch.multiprocessing.Queue,
+):
+    """Worker function for multi-branch CFG parallel test."""
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
+
+    update_environment_variables(
+        {
+            "RANK": str(local_rank),
+            "LOCAL_RANK": str(local_rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": "29504",
+        }
+    )
+
+    init_distributed_environment()
+    initialize_model_parallel(cfg_parallel_size=cfg_parallel_size)
+
+    cfg_rank = get_classifier_free_guidance_rank()
+    cfg_world_size = get_classifier_free_guidance_world_size()
+    assert cfg_world_size == cfg_parallel_size
+
+    pipeline = MultiBranchTestPipeline(
+        in_channels=test_config["channels"],
+        hidden_dim=test_config["hidden_dim"],
+        seed=test_config["model_seed"],
+    )
+    pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype)
+    pipeline.transformer.eval()
+
+    n_branches = test_config["n_branches"]
+    batch_size = test_config["batch_size"]
+    channels = test_config["channels"]
+    height = test_config["height"]
+    width = test_config["width"]
+
+    # Create N branch inputs with distinct seeds
+    branches_kwargs = []
+    for b in range(n_branches):
+        torch.manual_seed(test_config["input_seed"] + b)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(test_config["input_seed"] + b)
+        x = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device)
+        branches_kwargs.append({"x": x})
+
+    with torch.no_grad():
+        noise_pred = pipeline.predict_noise_with_multi_branch_cfg(
+            do_true_cfg=True,
+            true_cfg_scale=test_config["cfg_scale"],
+            branches_kwargs=branches_kwargs,
+            cfg_normalize=test_config["cfg_normalize"],
+        )
+
+    assert noise_pred is not None
+    result_queue.put((cfg_rank, noise_pred.cpu()))
+
+    destroy_distributed_env()
+
+
+def _test_multi_branch_sequential_worker(
+    local_rank: int,
+    world_size: int,
+    dtype: torch.dtype,
+    test_config: dict,
+    result_queue: torch.multiprocessing.Queue,
+):
+    """Worker function for sequential multi-branch CFG test (baseline)."""
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
+
+    update_environment_variables(
+        {
+            "RANK": str(local_rank),
+            "LOCAL_RANK": str(local_rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": "29505",
+        }
+    )
+
+    init_distributed_environment()
+    initialize_model_parallel(cfg_parallel_size=1)
+
+    cfg_world_size = get_classifier_free_guidance_world_size()
+    assert cfg_world_size == 1
+
+    pipeline = MultiBranchTestPipeline(
+        in_channels=test_config["channels"],
+        hidden_dim=test_config["hidden_dim"],
+        seed=test_config["model_seed"],
+    )
+    pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype)
+    pipeline.transformer.eval()
+
+    n_branches = test_config["n_branches"]
+    batch_size = test_config["batch_size"]
+    channels = test_config["channels"]
+    height = test_config["height"]
+    width = test_config["width"]
+
+    branches_kwargs = []
+    for b in range(n_branches):
+        torch.manual_seed(test_config["input_seed"] + b)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(test_config["input_seed"] + b)
+        x = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device)
+        branches_kwargs.append({"x": x})
+
+    with torch.no_grad():
+        noise_pred = pipeline.predict_noise_with_multi_branch_cfg(
+            do_true_cfg=True,
+            true_cfg_scale=test_config["cfg_scale"],
+            branches_kwargs=branches_kwargs,
+            cfg_normalize=test_config["cfg_normalize"],
+        )
+
+    assert noise_pred is not None
+    result_queue.put(noise_pred.cpu())
+
+    destroy_distributed_env()
+
+
+@pytest.mark.parametrize(
+    "cfg_parallel_size,n_branches",
+    [
+        (2, 2),  # 2 branches on 2 GPUs: [[0],[1]]
+        (2, 3),  # 3 branches on 2 GPUs: [[0,2],[1]]
+        (3, 3),  # 3 branches on 3 GPUs: [[0],[1],[2]]
+        (2, 4),  # 4 branches on 2 GPUs: [[0,2],[1,3]]
+    ],
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("batch_size", [2])
+@pytest.mark.parametrize("cfg_normalize", [False, True])
+def test_predict_noise_with_multi_branch_cfg(
+    cfg_parallel_size: int,
+    n_branches: int,
+    dtype: torch.dtype,
+    batch_size: int,
+    cfg_normalize: bool,
+):
+    """
+    Test that predict_noise_with_multi_branch_cfg produces identical results
+    with and without CFG parallel for N-branch models.
+
+    Args:
+        cfg_parallel_size: Number of GPUs for CFG parallel
+        n_branches: Number of CFG branches
+        dtype: Data type for computation
+        batch_size: Batch size for testing
+        cfg_normalize: Whether to normalize CFG output
+    """
+    available_gpus = current_omni_platform.get_device_count()
+    if available_gpus < cfg_parallel_size:
+        pytest.skip(f"Test requires {cfg_parallel_size} GPUs but only {available_gpus} available")
+
+    if n_branches == 2:
+        cfg_scale = 5.0
+    elif n_branches == 3:
+        cfg_scale = {"text": 5.0, "image": 2.0}
+    else:
+        cfg_scale = {"text": 5.0, "image": 2.0, "vid_ref": 1.5}
+
+    test_config = {
+        "batch_size": batch_size,
+        "channels": 4,
+        "height": 16,
+        "width": 16,
+        "hidden_dim": 128,
+        "cfg_scale": cfg_scale,
+        "cfg_normalize": cfg_normalize,
+        "model_seed": 42,
+        "input_seed": 123,
+        "n_branches": n_branches,
+    }
+
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    baseline_queue = manager.Queue()
+    cfg_parallel_queue = manager.Queue()
+
+    # Run baseline (sequential, cfgp=1)
+    torch.multiprocessing.spawn(
+        _test_multi_branch_sequential_worker,
+        args=(1, dtype, test_config, baseline_queue),
+        nprocs=1,
+    )
+
+    # Run CFG parallel
+    torch.multiprocessing.spawn(
+        _test_multi_branch_parallel_worker,
+        args=(cfg_parallel_size, cfg_parallel_size, dtype, test_config, cfg_parallel_queue),
+        nprocs=cfg_parallel_size,
+    )
+
+    baseline_output = baseline_queue.get()
+    cfg_parallel_outputs = [cfg_parallel_queue.get() for _ in range(cfg_parallel_size)]
+    cfg_parallel_outputs.sort(key=lambda item: item[0])
+    cfg_parallel_output = cfg_parallel_outputs[0][1]
+
+    # All ranks should produce identical output
+    for cfg_rank, rank_output in cfg_parallel_outputs[1:]:
+        torch.testing.assert_close(
+            rank_output,
+            cfg_parallel_output,
+            rtol=0,
+            atol=0,
+            msg=f"Multi-branch CFG parallel ranks differ (rank 0 vs rank {cfg_rank})",
+        )
+
+    assert baseline_output.shape == cfg_parallel_output.shape, (
+        f"Shape mismatch: baseline {baseline_output.shape} vs CFG parallel {cfg_parallel_output.shape}"
+    )
+
+    if dtype == torch.float32:
+        rtol, atol = 1e-5, 1e-5
+    elif dtype == torch.bfloat16:
+        rtol, atol = 1e-2, 1e-2
+    else:
+        rtol, atol = 1e-3, 1e-3
+
+    torch.testing.assert_close(
+        cfg_parallel_output,
+        baseline_output,
+        rtol=rtol,
+        atol=atol,
+        msg=(
+            f"Multi-branch CFG parallel output differs from sequential\n"
+            f"  n_branches={n_branches}, cfg_parallel_size={cfg_parallel_size}\n"
+            f"  dtype={dtype}, cfg_normalize={cfg_normalize}\n"
+            f"  Max diff: {(cfg_parallel_output - baseline_output).abs().max().item():.6e}"
+        ),
+    )
+
+    print(
+        f"✓ Test passed: multi_branch n_branches={n_branches}, "
+        f"cfg_size={cfg_parallel_size}, dtype={dtype}, cfg_normalize={cfg_normalize}"
+    )
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+def test_multi_branch_without_cfg(dtype: torch.dtype):
+    """
+    Test predict_noise_with_multi_branch_cfg when do_true_cfg=False.
+
+    When CFG is disabled, only the first branch (positive) should be computed.
+    This test runs on a single GPU without distributed environment.
+    """
+    available_gpus = current_omni_platform.get_device_count()
+    if available_gpus < 1:
+        pytest.skip("Test requires at least 1 GPU")
+
+    device = torch.device(f"{current_omni_platform.device_type}:0")
+    current_omni_platform.set_device(device)
+
+    pipeline = MultiBranchTestPipeline(in_channels=4, hidden_dim=128, seed=42)
+    pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype)
+    pipeline.transformer.eval()
+
+    # Create 3 branch inputs (only first should be used)
+    branches_kwargs = []
+    for b in range(3):
+        torch.manual_seed(123 + b)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(123 + b)
+        x = torch.randn(1, 4, 16, 16, dtype=dtype, device=device)
+        branches_kwargs.append({"x": x})
+
+    with torch.no_grad():
+        noise_pred = pipeline.predict_noise_with_multi_branch_cfg(
+            do_true_cfg=False,  # No CFG
+            true_cfg_scale=5.0,
+            branches_kwargs=branches_kwargs,
+            cfg_normalize=False,
+        )
+
+    assert noise_pred is not None
+    assert noise_pred.shape == (1, 4, 16, 16)
+
+    print(f"✓ Test passed: multi_branch predict_noise without CFG (dtype={dtype})")

vllm_omni/diffusion/distributed/cfg_parallel.py

@@ -9,6 +9,7 @@ from abc import ABCMeta
 from typing import Any
 
 import torch
+from vllm.logger import init_logger
 
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_cfg_group,
@@ -16,6 +17,8 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     get_classifier_free_guidance_world_size,
 )
 
+logger = init_logger(__name__)
+
 
 def _wrap(pred: torch.Tensor | tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
     """Normalize prediction to tuple form."""
@@ -32,6 +35,24 @@ def _slice_pred(pred: tuple[torch.Tensor, ...], output_slice: int) -> tuple[torc
     return tuple(p[:, :output_slice] for p in pred)
 
 
+def _dispatch_branches(n_branches: int, n_ranks: int) -> list[list[int]]:
+    """
+    Round-robin dispatch N branches to M ranks.
+
+    Rule: branch i → rank (i % n_ranks).
+
+    Examples:
+        _dispatch_branches(3, 2) -> [[0, 2], [1]]
+        _dispatch_branches(3, 3) -> [[0], [1], [2]]
+        _dispatch_branches(4, 2) -> [[0, 2], [1, 3]]
+        _dispatch_branches(4, 4) -> [[0], [1], [2], [3]]
+    """
+    assignments: list[list[int]] = [[] for _ in range(n_ranks)]
+    for i in range(n_branches):
+        assignments[i % n_ranks].append(i)
+    return assignments
+
+
 class CFGParallelMixin(metaclass=ABCMeta):
     """
     Base Mixin class for Diffusion pipelines providing shared CFG methods.
@@ -189,6 +210,165 @@ class CFGParallelMixin(metaclass=ABCMeta):
             results.append(comb)
         return _unwrap(tuple(results))
 
+    # ── N-branch CFG interface (for 3+ branch models) ──
+
+    def predict_noise_with_multi_branch_cfg(
+        self,
+        do_true_cfg: bool,
+        true_cfg_scale: float | dict[str, float],
+        branches_kwargs: list[dict[str, Any]],
+        cfg_normalize: bool = False,
+        output_slice: int | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+        """
+        Predict noise with N-branch CFG dispatch across M GPUs.
+
+        This is the multi-branch counterpart of predict_noise_maybe_with_cfg().
+        Use this for models with 3 or more CFG branches (e.g., OmniGen2, Bagel,
+        DreamID). Existing 2-branch models should continue using
+        predict_noise_maybe_with_cfg().
+
+        Args:
+            do_true_cfg: Whether to apply CFG.
+            true_cfg_scale: CFG scale factor (passed to combine_multi_branch_cfg_noise).
+            branches_kwargs: List of N dicts, each containing kwargs for one
+                predict_noise() call. branches_kwargs[0] is always the
+                positive/conditional branch.
+            cfg_normalize: Whether to normalize (passed to combine_multi_branch_cfg_noise).
+            output_slice: If set, slice each output to [:, :output_slice].
+
+        Returns:
+            Combined noise prediction, identical on all ranks in CFG parallel.
+        """
+        if do_true_cfg:
+            n_branches = len(branches_kwargs)
+            cfg_world_size = get_classifier_free_guidance_world_size()
+            cfg_parallel_ready = cfg_world_size > 1
+
+            if cfg_parallel_ready:
+                return self._predict_multi_branch_parallel(
+                    branches_kwargs,
+                    n_branches,
+                    cfg_world_size,
+                    true_cfg_scale,
+                    cfg_normalize,
+                    output_slice,
+                )
+            else:
+                # Sequential: run all N branches on single device
+                preds: list[torch.Tensor | tuple[torch.Tensor, ...]] = []
+                for kw in branches_kwargs:
+                    pred = _wrap(self.predict_noise(**kw))
+                    if output_slice is not None:
+                        pred = _slice_pred(pred, output_slice)
+                    preds.append(_unwrap(pred))
+                return self.combine_multi_branch_cfg_noise(preds, true_cfg_scale, cfg_normalize)
+        else:
+            # No CFG: only compute positive/conditional prediction
+            pred = self.predict_noise(**branches_kwargs[0])
+            if output_slice is not None:
+                pred = _unwrap(_slice_pred(_wrap(pred), output_slice))
+            return pred
+
+    def _predict_multi_branch_parallel(
+        self,
+        branches_kwargs: list[dict[str, Any]],
+        n_branches: int,
+        cfg_world_size: int,
+        true_cfg_scale: float,
+        cfg_normalize: bool,
+        output_slice: int | None,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+        """Dispatch N branches across M ranks, all_gather, then combine."""
+        cfg_group = get_cfg_group()
+        cfg_rank = get_classifier_free_guidance_rank()
+
+        if cfg_world_size > n_branches:
+            logger.warning_once(
+                "cfg_parallel_size=%d > n_branches=%d, %d GPU(s) will be idle for CFG",
+                cfg_world_size,
+                n_branches,
+                cfg_world_size - n_branches,
+            )
+
+        # Assign branches to ranks via round-robin
+        assignments = _dispatch_branches(n_branches, cfg_world_size)
+        my_branch_ids = assignments[cfg_rank]
+        max_per_rank = max(len(a) for a in assignments)
+
+        # Run assigned branches
+        my_preds: list[tuple[torch.Tensor, ...]] = []
+        for bid in my_branch_ids:
+            pred = _wrap(self.predict_noise(**branches_kwargs[bid]))
+            if output_slice is not None:
+                pred = _slice_pred(pred, output_slice)
+            my_preds.append(pred)
+
+        # Idle ranks (cfg_world_size > n_branches) run a forward pass to get the output shape for all_gather.
+        # Output shape cannot be inferred from kwargs — may be tuple, sliced, etc.
+        if not my_preds:
+            pred = _wrap(self.predict_noise(**branches_kwargs[0]))
+            if output_slice is not None:
+                pred = _slice_pred(pred, output_slice)
+            my_preds.append(pred)
+
+        # Pad to max_per_rank with zeros so all ranks have same size
+        ref_pred = my_preds[0]
+        while len(my_preds) < max_per_rank:
+            my_preds.append(tuple(torch.zeros_like(t) for t in ref_pred))
+
+        # All-gather each output element separately (like predict_noise_maybe_with_cfg)
+        # For each slot, gather across ranks; then pick valid results by owner_rank
+        # all_slots[slot][elem_idx] = [rank0_tensor, rank1_tensor, ...]
+        all_slots: list[list[list[torch.Tensor]]] = []
+        for slot in range(max_per_rank):
+            slot_results: list[list[torch.Tensor]] = []
+            for p in my_preds[slot]:
+                gathered = cfg_group.all_gather(p, separate_tensors=True)
+                slot_results.append(gathered)
+            all_slots.append(slot_results)
+
+        # Reconstruct final_preds in branch order
+        final_preds: list[torch.Tensor | tuple[torch.Tensor, ...]] = []
+        for bid in range(n_branches):
+            owner_rank = bid % cfg_world_size
+            slot_idx = bid // cfg_world_size
+            elements = tuple(all_slots[slot_idx][elem_idx][owner_rank] for elem_idx in range(len(ref_pred)))
+            final_preds.append(_unwrap(elements))
+
+        return self.combine_multi_branch_cfg_noise(final_preds, true_cfg_scale, cfg_normalize)
+
+    def combine_multi_branch_cfg_noise(
+        self,
+        predictions: list[torch.Tensor | tuple[torch.Tensor, ...]],
+        true_cfg_scale: float | dict[str, float],
+        cfg_normalize: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+        """
+        Combine N branch predictions. Default: standard 2-branch CFG formula.
+
+        Override this method for custom multi-branch combine logic.
+
+        Args:
+            predictions: List of N predictions, where predictions[0] is always
+                the positive/conditional branch.
+            true_cfg_scale: CFG scale factor (float for 2-branch, dict for multi-branch).
+            cfg_normalize: Whether to normalize the combined prediction.
+
+        Returns:
+            Combined noise prediction.
+        """
+        positive = _wrap(predictions[0])
+        negative = _wrap(predictions[1])
+
+        results = []
+        for p, n in zip(positive, negative):
+            comb = n + true_cfg_scale * (p - n)
+            if cfg_normalize:
+                comb = self.cfg_normalize_function(p, comb)
+            results.append(comb)
+        return _unwrap(tuple(results))
+
     def predict_noise(self, *args: Any, **kwargs: Any) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """
         Forward pass through transformer to predict noise.

vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py

@@ -15,11 +15,6 @@ from tqdm import tqdm
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
-from vllm_omni.diffusion.distributed.parallel_state import (
-    get_cfg_group,
-    get_classifier_free_guidance_rank,
-    get_classifier_free_guidance_world_size,
-)
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportImageInput
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -249,6 +244,28 @@ class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Suppor
 
         return sample_scheduler, timesteps
 
+    def predict_noise(self, **kwargs):
+        pred_vid, pred_audio = self.model(**kwargs)
+        return (pred_vid[0], pred_audio[0])
+
+    def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize=False):
+        vid_pos, audio_pos = predictions[0]
+        vid_neg, audio_neg = predictions[1]
+        vid_ip_neg, _ = predictions[2]
+        _, refaudio_neg = predictions[3]
+
+        pred_video = (
+            vid_neg
+            + true_cfg_scale["video_cfg_scale"] * (vid_pos - vid_neg)
+            + true_cfg_scale["video_ref_cfg_scale"] * (vid_pos - vid_ip_neg)
+        )
+        pred_audio = (
+            audio_neg
+            + true_cfg_scale["audio_cfg_scale"] * (audio_pos - audio_neg)
+            + true_cfg_scale["audio_ref_cfg_scale"] * (audio_pos - refaudio_neg)
+        )
+        return (pred_video, pred_audio)
+
     def diffuse(
         self,
         video_noise: torch.Tensor,
@@ -306,72 +323,22 @@ class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Suppor
                 "vid_context": [text_embeddings_video_neg],
             }
 
-            if get_classifier_free_guidance_world_size() > 1:
-                # Enable CFG-parallel: rank0 computes positive, rank1 computes negative.
-                cfg_group = get_cfg_group()
-                cfg_rank = get_classifier_free_guidance_rank()
+            branches_kwargs = [
+                {"vid": [model_input_video], "audio": [model_input_audio], "t": timestep_input, **pos_args},
+                {"vid": [model_input_video], "audio": [model_input_audio], "t": timestep_input, **neg_args},
+                {"vid": [model_input_video_neg], "audio": [model_input_audio], "t": timestep_input, **pos_args},
+                {"vid": [model_input_video], "audio": [model_input_audio_neg], "t": timestep_input, **pos_args},
+            ]
 
-                if cfg_rank == 0:
-                    pred_vid, pred_audio = self.model(
-                        vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **pos_args
-                    )
-                    pre_vid_ip_neg, _ = self.model(
-                        vid=[model_input_video_neg], audio=[model_input_audio], t=timestep_input, **pos_args
-                    )
-                    pred_vid_0 = pred_vid[0]
-                    pred_audio_0 = pred_audio[0]
-                    pre_vid_ip_0 = pre_vid_ip_neg[0]
-                    pred_refaudio_0 = torch.zeros_like(pred_audio_0)  # dummy tensor
-                else:
-                    pred_vid, pred_audio = self.model(
-                        vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **neg_args
-                    )
-                    _, pred_refaudio_neg = self.model(
-                        vid=[model_input_video], audio=[model_input_audio_neg], t=timestep_input, **pos_args
-                    )
-                    pred_vid_0 = pred_vid[0]
-                    pred_audio_0 = pred_audio[0]
-                    pre_vid_ip_0 = torch.zeros_like(pred_vid_0)  # dummy tensor
-                    pred_refaudio_0 = pred_refaudio_neg[0]
-
-                pred_vid_gathered = cfg_group.all_gather(pred_vid_0, separate_tensors=True)
-                pred_audio_gathered = cfg_group.all_gather(pred_audio_0, separate_tensors=True)
-                pre_vid_ip_gathered = cfg_group.all_gather(pre_vid_ip_0, separate_tensors=True)
-                pred_refaudio_gathered = cfg_group.all_gather(pred_refaudio_0, separate_tensors=True)
-
-                pred_vid_pos = [pred_vid_gathered[0]]
-                pred_vid_neg = [pred_vid_gathered[1]]
-                pred_audio_pos = [pred_audio_gathered[0]]
-                pred_audio_neg = [pred_audio_gathered[1]]
-                pre_vid_ip_neg = [pre_vid_ip_gathered[0]]
-                pred_refaudio_neg = [pred_refaudio_gathered[1]]
-            else:
-                pred_vid_pos, pred_audio_pos = self.model(
-                    vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **pos_args
-                )
-
-                pred_vid_neg, pred_audio_neg = self.model(
-                    vid=[model_input_video], audio=[model_input_audio], t=timestep_input, **neg_args
-                )
-
-                pre_vid_ip_neg, _ = self.model(
-                    vid=[model_input_video_neg], audio=[model_input_audio], t=timestep_input, **pos_args
-                )
-
-                _, pred_refaudio_neg = self.model(
-                    vid=[model_input_video], audio=[model_input_audio_neg], t=timestep_input, **pos_args
-                )
-
-            pred_video_guided = (
-                pred_vid_neg[0]
-                + self.video_cfg_scale * (pred_vid_pos[0] - pred_vid_neg[0])
-                + self.video_ref_cfg_scale * (pred_vid_pos[0] - pre_vid_ip_neg[0])
-            )
-
-            pred_audio_guided = (
-                pred_audio_neg[0]
-                + self.audio_cfg_scale * (pred_audio_pos[0] - pred_audio_neg[0])
-                + self.audio_ref_cfg_scale * (pred_audio_pos[0] - pred_refaudio_neg[0])
+            pred_video_guided, pred_audio_guided = self.predict_noise_with_multi_branch_cfg(
+                do_true_cfg=True,
+                true_cfg_scale={
+                    "video_cfg_scale": self.video_cfg_scale,
+                    "video_ref_cfg_scale": self.video_ref_cfg_scale,
+                    "audio_cfg_scale": self.audio_cfg_scale,
+                    "audio_ref_cfg_scale": self.audio_ref_cfg_scale,
+                },
+                branches_kwargs=branches_kwargs,
             )
             video_noise = scheduler_video.step(
                 pred_video_guided.unsqueeze(0), t_v, video_noise.unsqueeze(0), return_dict=False

vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py

@@ -29,6 +29,7 @@ from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.omnigen2.omnigen2_transformer import (
@@ -619,7 +620,7 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-class OmniGen2Pipeline(nn.Module):
+class OmniGen2Pipeline(CFGParallelMixin, nn.Module):
     """
     Pipeline for text-to-image generation using OmniGen2.
 
@@ -1171,14 +1172,6 @@ class OmniGen2Pipeline(nn.Module):
         self._num_timesteps = len(timesteps)
 
         for i, t in enumerate(timesteps):
-            model_pred = self.predict(
-                t=t,
-                latents=latents,
-                prompt_embeds=prompt_embeds,
-                freqs_cis=freqs_cis,
-                prompt_attention_mask=prompt_attention_mask,
-                ref_image_hidden_states=ref_latents,
-            )
             text_guidance_scale = (
                 self.text_guidance_scale if self.cfg_range[0] <= i / len(timesteps) <= self.cfg_range[1] else 1.0
             )
@@ -1186,8 +1179,26 @@ class OmniGen2Pipeline(nn.Module):
                 self.image_guidance_scale if self.cfg_range[0] <= i / len(timesteps) <= self.cfg_range[1] else 1.0
             )
 
+            positive_kwargs = dict(
+                t=t,
+                latents=latents,
+                prompt_embeds=prompt_embeds,
+                freqs_cis=freqs_cis,
+                prompt_attention_mask=prompt_attention_mask,
+                ref_image_hidden_states=ref_latents,
+            )
+            uncond_kwargs = dict(
+                t=t,
+                latents=latents,
+                prompt_embeds=negative_prompt_embeds,
+                freqs_cis=freqs_cis,
+                prompt_attention_mask=negative_prompt_attention_mask,
+                ref_image_hidden_states=None,
+            )
+
             if text_guidance_scale > 1.0 and image_guidance_scale > 1.0:
-                model_pred_ref = self.predict(
+                # 3-branch CFG: pos + ref_neg + uncond
+                ref_neg_kwargs = dict(
                     t=t,
                     latents=latents,
                     prompt_embeds=negative_prompt_embeds,
@@ -1195,31 +1206,24 @@ class OmniGen2Pipeline(nn.Module):
                     prompt_attention_mask=negative_prompt_attention_mask,
                     ref_image_hidden_states=ref_latents,
                 )
-
-                model_pred_uncond = self.predict(
-                    t=t,
-                    latents=latents,
-                    prompt_embeds=negative_prompt_embeds,
-                    freqs_cis=freqs_cis,
-                    prompt_attention_mask=negative_prompt_attention_mask,
-                    ref_image_hidden_states=None,
-                )
-
-                model_pred = (
-                    model_pred_uncond
-                    + image_guidance_scale * (model_pred_ref - model_pred_uncond)
-                    + text_guidance_scale * (model_pred - model_pred_ref)
+                model_pred = self.predict_noise_with_multi_branch_cfg(
+                    do_true_cfg=True,
+                    true_cfg_scale={
+                        "text": text_guidance_scale,
+                        "image": image_guidance_scale,
+                    },
+                    branches_kwargs=[positive_kwargs, ref_neg_kwargs, uncond_kwargs],
                 )
             elif text_guidance_scale > 1.0:
-                model_pred_uncond = self.predict(
-                    t=t,
-                    latents=latents,
-                    prompt_embeds=negative_prompt_embeds,
-                    freqs_cis=freqs_cis,
-                    prompt_attention_mask=negative_prompt_attention_mask,
-                    ref_image_hidden_states=None,
+                # 2-branch CFG: pos + uncond
+                model_pred = self.predict_noise_with_multi_branch_cfg(
+                    do_true_cfg=True,
+                    true_cfg_scale=text_guidance_scale,
+                    branches_kwargs=[positive_kwargs, uncond_kwargs],
                 )
-                model_pred = model_pred_uncond + text_guidance_scale * (model_pred - model_pred_uncond)
+            else:
+                # No CFG
+                model_pred = self.predict_noise(**positive_kwargs)
 
             latents = self.scheduler.step(model_pred, t, latents, return_dict=False)[0]
 
@@ -1265,6 +1269,21 @@ class OmniGen2Pipeline(nn.Module):
         )
         return model_pred
 
+    def predict_noise(self, **kwargs):
+        """Override CFGParallelMixin.predict_noise to use self.predict."""
+        return self.predict(**kwargs)
+
+    def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize=False):
+        """Override: 3-branch dual scale or 2-branch standard CFG."""
+        if len(predictions) == 3:
+            text_scale = true_cfg_scale["text"]
+            image_scale = true_cfg_scale["image"]
+            pos, ref, uncond = predictions[0], predictions[1], predictions[2]
+            return uncond + image_scale * (ref - uncond) + text_scale * (pos - ref)
+        # 2-branch: standard CFG
+        pos, neg = predictions[0], predictions[1]
+        return neg + true_cfg_scale * (pos - neg)
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)