This section describes how to add CFG-Parallel (Classifier-Free Guidance Parallel) to a diffusion pipeline. We use the Qwen-Image pipeline as the reference implementation.
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.
vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic. Pipelines inherit from this mixin and implement a `diffuse()` method that orchestrates the denoising loop.
| [`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 |
- 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.
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.
# Run diffusion loop (calls the mixin's diffuse method)
latents=self.diffuse(
prompt_embeds=prompt_embeds,
prompt_embeds_mask=prompt_embeds_mask,
negative_prompt_embeds=negative_embeds,
negative_prompt_embeds_mask=negative_mask,
latents=latents,
timesteps=timesteps,
do_true_cfg=do_true_cfg,
true_cfg_scale=guidance_scale,
...
)
```
---
## Customization
### Override `predict_noise()` for Custom Transformer Calls
If your transformer requires custom prediction function, you can rewrite `predict_noise` function. Taking Wan2.2 as an example, which has two transformer models. The actual transformer to be called is determined by `self.transformer`.
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`.
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
When each output has its own denoising schedule, implement a composite scheduler that dispatches to per-output schedulers. Assign it to `self.scheduler` so the default `scheduler_step()` works without override.
**Complete example (video + audio with separate schedulers and diffuse loop):**
```python
classVideoAudioScheduler:
"""Composite scheduler dispatching to video and audio schedulers."""
> **Note:** If you use a non-deterministic scheduler, e.g., DDPM, please set `self.scheduler_step_maybe_with_cfg(..., generator=torch.Generator(device).manual_seed(seed))` explicitly to control the randomness of scheduler step among ranks.
