[Feature][Diffusion]: CFG parallelism support for Boogu-Image (#6786)

Signed-off-by: nagisa <1434936049@qq.com> Signed-off-by: nagisa-kun <1434936049@qq.com>
2026-09-03 10:37:21 +08:00
1 parent 376b28aced
commit 8f284b343a
@@ -229,6 +229,17 @@ steps:
            --run-level "full_model"
        mirror_hardwares: h100_2

+      - label: ":full_moon: Diffusion X2I(&A&T) · Boogu Edit CFG · 3-GPU"
+        timeout_in_minutes: 120
+        commands:
+          - >-
+            pytest -sv
+            tests/e2e/online_serving/test_boogu_image_edit_expansion.py
+            -k "test_double_guidance_cfg_parallel and cfg_parallel_3_full"
+            -m "slow and diffusion and cuda and H100 and cards_3"
+            --run-level "full_model"
+        mirror_hardwares: h100_3
+
      - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4 · 4-GPU"
        timeout_in_minutes: 120
        commands:
@@ -94,7 +94,7 @@ In `DiffusionParallelConfig`


 !!! info
-    Most models support `cfg_parallel_size=2` (positive branch on rank 0, negative branch on rank 1). **Bagel** is an exception: it supports `cfg_parallel_size=3`, which adds a third branch on rank 2 for full three-way CFG parallelism.
+    Most models support `cfg_parallel_size=2` (positive branch on rank 0, negative branch on rank 1). **Bagel** and **Boogu-Image-Edit double guidance** also support `cfg_parallel_size=3` for full three-way branch execution. Boogu double guidance additionally supports size 2 via round-robin assignment of its three branches.


 ---
@@ -117,6 +117,7 @@ The following tables show which models support each feature:
 | Model                    | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
 |--------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
 | **Bagel**                |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **Boogu-Image**          |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **FLUX.1-dev**           |     ✅     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **FLUX.1-schnell**       |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **FLUX.2-klein**         |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
@@ -39,9 +39,10 @@ rendering. The Edit and Edit-Turbo checkpoints use the same native

 ## Hardware Support

-This recipe documents tested configurations for CUDA GPU serving. The native
-pipeline runs single-GPU; multi-GPU parallelism, CPU offload, and cache
-acceleration are not yet supported for this model (see Notes).
+This recipe documents supported configurations for CUDA GPU serving. The native
+pipeline supports single-GPU inference and classifier-free-guidance (CFG)
+parallelism. CPU offload, cache acceleration, and other multi-GPU dimensions
+remain unsupported (see Notes).

 ## GPU

@@ -108,8 +109,24 @@ curl -s http://localhost:8091/v1/chat/completions \
  The model's maximum native resolution is 2K.
 - **Known limitations (not yet supported):** CPU offload
  (`--enable-cpu-offload` / `--enable-layerwise-offload`), Cache-DiT
-  (`--cache-backend cache_dit`), and multi-GPU parallelism (TP / SP / CFG /
-  HSDP) are planned follow-ups and are not validated for this model yet.
+  (`--cache-backend cache_dit`), and TP / SP / HSDP multi-GPU parallelism.
+
+### 2 x H100 (CFG parallel, Base T2I)
+
+Base T2I has two predictions per guided denoising step, so two CFG ranks map
+one-to-one to the positive and negative-text branches:
+
+```bash
+vllm serve Boogu/Boogu-Image-0.1-Base \
+  --omni \
+  --cfg-parallel-size 2 \
+  --port 8091
+```
+
+Use `guidance_scale > 1.0` to activate the parallel two-branch path. Requests
+with `guidance_scale=1.0` remain valid on the same server: every rank evaluates
+only the positive branch, no negative embeddings are built, and no guidance is
+applied.

 ## Image editing (Boogu-Image-0.1-Edit)

@@ -124,6 +141,33 @@ is sent).
 vllm serve Boogu/Boogu-Image-0.1-Edit --omni --port 8091
 ```

+For text-only or image-only guidance, use two CFG ranks:
+
+```bash
+vllm serve Boogu/Boogu-Image-0.1-Edit \
+  --omni \
+  --cfg-parallel-size 2 \
+  --port 8091
+```
+
+Double guidance has three predictions per denoising step. It supports either
+two ranks (round-robin assignment: rank 0 runs branches 0 and 2) or three ranks
+(one branch per rank):
+
+```bash
+# Practical two-GPU configuration
+vllm serve Boogu/Boogu-Image-0.1-Edit \
+  --omni \
+  --cfg-parallel-size 2 \
+  --port 8091
+
+# Maximum branch parallelism
+vllm serve Boogu/Boogu-Image-0.1-Edit \
+  --omni \
+  --cfg-parallel-size 3 \
+  --port 8091
+```
+
 ### Verification

 Edit an image with `/v1/images/edits` (the model-card example — change a photo
@@ -175,11 +219,29 @@ curl -s http://localhost:8091/v1/chat/completions \
    default `1.0` = off). Setting it `> 1.0` enables the double-guidance path
    (3 model predictions per step), steering more strongly toward the reference
    image.
+- **CFG-size recommendations:**
+
+  | Request mode | Scales | Predictions per step | Recommended CFG size |
+  | --- | --- | ---: | ---: |
+  | Base T2I text guidance | `guidance_scale > 1`, no reference | 2 | 2 |
+  | Edit text-only guidance | `guidance_scale > 1`, `guidance_scale_2 = 1` | 2 | 2 |
+  | Edit image-only guidance | `guidance_scale = 1`, `guidance_scale_2 > 1` | 2 | 2 |
+  | Edit double guidance | both scales `> 1` | 3 | 2 (practical) or 3 (maximum parallelism) |
+  | CFG off | both scales `= 1` | 1 | Reuse a size-2 server safely; no speedup expected |
+
+  Double guidance preserves Boogu's original combination:
+
+  ```text
+  pred = positive_with_reference
+       + (text_scale - 1) * (positive_with_reference - negative_with_reference)
+       + (image_scale - 1) * (negative_with_reference - uncond)
+  ```
 - **Output resolution:** the output size follows the reference image (upstream
  `align_res`, on by default for a single reference), so `height`/`width` are
  derived from the input and requested sizes are not applied for edits.
- **Same limitations** as the Base checkpoint apply (no CPU offload, Cache-DiT,
-  or multi-GPU parallelism yet).
+- **Same limitations** as the Base checkpoint apply except that CFG parallelism
+  is supported as described above; CPU offload, Cache-DiT, TP, SP, and HSDP are
+  still unsupported.
 - **Offline editing:** the shared offline example
  [`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_to_image.md)
  supports Boogu-Image-Edit directly
@@ -251,3 +313,21 @@ curl -s http://localhost:8091/v1/chat/completions \
  the regular Edit checkpoint applies.
 - **Known limitations:** the same single-GPU limitations as Base and Edit apply;
  CPU offload, Cache-DiT, and multi-GPU parallelism are not yet validated.
+
+## Performance validation
+
+The checked-in single-device A40 measurements remain the current Boogu
+baseline: Base T2I is about 16.37 s request latency at concurrency 1 and Edit
+text-only guidance is about 29.64 s for 512x512, 28-step requests, with roughly
+36.8 GB peak memory. CFG-parallel A/B numbers must be collected on the same
+hardware and software revision before reporting a speedup; do not compare a
+multi-GPU result against these A40 numbers as if it were a controlled A/B.
+
+Run the dedicated perf rows in
+`tests/dfx/perf/tests/test_boogu_image_vllm_omni.json` and
+`tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json` on 2x/3x H100, then
+record latency, throughput, and per-GPU peak memory beside the single-device
+run. The expected shape is close to the general CFG-parallel range (roughly
+1.5x-1.8x for balanced two-branch workloads); the three-branch Edit result will
+depend on whether two-rank round-robin or full three-rank execution is used and
+must be reported from measurement rather than assumed.
@@ -1,7 +1,7 @@
 [
    {
        "test_name": "test_boogu_image_edit_single_device",
-        "description": "Single-device baseline (no parallelism), concurrency sweep 1/2/4 at 512x512, 28 steps. Text-only edit path (2 predictions/step); a double-guidance (guidance_scale_2) case is a follow-up since the random dataset cannot inject extra_body params. Baselines recorded on 1x A40 46GB (measured c1/c2/c4: qps 0.0337/0.0336/0.0336, latency_mean 29.64/56.50/101.48 s, peak mem 36754 MB) with ~10% margin.",
+        "description": "Single-device baseline (no parallelism), concurrency sweep 1/2/4 at 512x512, 28 steps. This row measures text-only edit guidance (2 predictions/step); dedicated rows below use extra-body to measure double guidance. Baselines recorded on 1x A40 46GB (measured c1/c2/c4: qps 0.0337/0.0336/0.0336, latency_mean 29.64/56.50/101.48 s, peak mem 36754 MB) with ~10% margin.",
        "server_type": "vllm-omni",
        "server_params": {
            "model": "Boogu/Boogu-Image-0.1-Edit",
@@ -33,5 +33,184 @@
                }
            }
        ]
+    },
+    {
+        "test_name": "test_boogu_image_edit_double_guidance_single_device",
+        "mark": [
+            {
+                "hardware_marks": {
+                    "res": {
+                        "cuda": "H100"
+                    },
+                    "num_cards": 1
+                }
+            },
+            "full_model",
+            "diffusion"
+        ],
+        "description": "Single-device sequential three-branch double-guidance control row for direct A/B comparison with cfg_parallel_size=2 and 3 on identical H100 hardware.",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Boogu/Boogu-Image-0.1-Edit",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps28_i2i_double_cfg1",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 28,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "warmup-requests": 4,
+                "warmup-concurrency": 1,
+                "warmup-num-inference-steps": 28,
+                "enable-negative-prompt": true,
+                "extra-body": {
+                    "guidance_scale": 5.0,
+                    "guidance_scale_2": 2.0,
+                    "seed": 42
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_boogu_image_edit_text_cfg_parallel_2",
+        "mark": [
+            {
+                "hardware_marks": {
+                    "res": {
+                        "cuda": "H100"
+                    },
+                    "num_cards": 2
+                }
+            },
+            "full_model",
+            "diffusion"
+        ],
+        "description": "Two-device CFG-parallel Edit text-only guidance measurement row. Compare with the single-device row on identical H100 hardware and software.",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Boogu/Boogu-Image-0.1-Edit",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps28_i2i_text_cfg2",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 28,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "warmup-requests": 4,
+                "warmup-concurrency": 1,
+                "warmup-num-inference-steps": 28,
+                "enable-negative-prompt": true,
+                "extra-body": {
+                    "guidance_scale": 5.0,
+                    "guidance_scale_2": 1.0,
+                    "seed": 42
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_boogu_image_edit_double_cfg_parallel_2",
+        "mark": [
+            {
+                "hardware_marks": {
+                    "res": {
+                        "cuda": "H100"
+                    },
+                    "num_cards": 2
+                }
+            },
+            "full_model",
+            "diffusion"
+        ],
+        "description": "Two-device round-robin execution of the three Edit double-guidance branches. Populate A/B latency, throughput, and per-GPU memory after GPU validation.",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Boogu/Boogu-Image-0.1-Edit",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps28_i2i_double_cfg2",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 28,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "warmup-requests": 4,
+                "warmup-concurrency": 1,
+                "warmup-num-inference-steps": 28,
+                "enable-negative-prompt": true,
+                "extra-body": {
+                    "guidance_scale": 5.0,
+                    "guidance_scale_2": 2.0,
+                    "seed": 42
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_boogu_image_edit_double_cfg_parallel_3",
+        "mark": [
+            {
+                "hardware_marks": {
+                    "res": {
+                        "cuda": "H100"
+                    },
+                    "num_cards": 3
+                }
+            },
+            "full_model",
+            "diffusion"
+        ],
+        "description": "Full three-device execution of the three Edit double-guidance branches. Compare directly with sequential and cfg_parallel_size=2 on identical H100 hardware.",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Boogu/Boogu-Image-0.1-Edit",
+            "serve_args": {
+                "cfg-parallel-size": 3,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps28_i2i_double_cfg3",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 28,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "warmup-requests": 4,
+                "warmup-concurrency": 1,
+                "warmup-num-inference-steps": 28,
+                "enable-negative-prompt": true,
+                "extra-body": {
+                    "guidance_scale": 5.0,
+                    "guidance_scale_2": 2.0,
+                    "seed": 42
+                }
+            }
+        ]
    }
 ]
@@ -33,5 +33,49 @@
                }
            }
        ]
+    },
+    {
+        "test_name": "test_boogu_image_cfg_parallel_2",
+        "mark": [
+            {
+                "hardware_marks": {
+                    "res": {
+                        "cuda": "H100"
+                    },
+                    "num_cards": 2
+                }
+            },
+            "full_model",
+            "diffusion"
+        ],
+        "description": "Two-device CFG-parallel Base T2I measurement row. Run beside test_boogu_image_single_device on the same H100 revision and populate latency, throughput, speedup, and per-GPU peak-memory baselines from the resulting report.",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Boogu/Boogu-Image-0.1-Base",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps28_cfg2",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 28,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "warmup-requests": 4,
+                "warmup-concurrency": 1,
+                "warmup-num-inference-steps": 28,
+                "enable-negative-prompt": true,
+                "extra-body": {
+                    "guidance_scale": 4.0,
+                    "seed": 42
+                }
+            }
+        ]
    }
 ]
@@ -17,7 +17,12 @@ import torch

 from tests.helpers.mark import hardware_marks, hardware_test
 from vllm_omni.diffusion.distributed import parallel_state
-from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin, _dispatch_branches, _wrap
+from vllm_omni.diffusion.distributed.cfg_parallel import (
+    CFGParallelMixin,
+    _dispatch_branches,
+    _get_cfg_world_size_or_one,
+    _wrap,
+)
 from vllm_omni.diffusion.distributed.parallel_state import (
    destroy_distributed_env,
    get_classifier_free_guidance_rank,
@@ -393,6 +398,68 @@ def _build_pipeline_on_device(
 # ---------------------------------------------------------------------------


+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_cfg_world_size_defaults_only_when_group_is_uninitialized(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(parallel_state, "_CFG", None)
+
+    assert not parallel_state.is_cfg_group_initialized()
+    assert _get_cfg_world_size_or_one() == 1
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_cfg_world_size_does_not_swallow_initialized_group_errors(monkeypatch: pytest.MonkeyPatch):
+    class BrokenCfgGroup:
+        @property
+        def world_size(self) -> int:
+            raise AssertionError("CFG group is broken")
+
+    monkeypatch.setattr(parallel_state, "_CFG", BrokenCfgGroup())
+
+    assert parallel_state.is_cfg_group_initialized()
+    with pytest.raises(AssertionError, match="CFG group is broken"):
+        _get_cfg_world_size_or_one()
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_cfg_defaults_to_sequential_without_initialized_group():
+    pipeline = _build_pipeline_on_device(
+        TestCFGPipeline,
+        hidden_dim=32,
+        model_seed=42,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+    positive_kwargs, negative_kwargs = _make_two_branch_inputs(
+        batch_size=1,
+        channels=4,
+        height=4,
+        width=4,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        input_seed=123,
+    )
+
+    with torch.no_grad():
+        positive = pipeline.predict_noise(**positive_kwargs)
+        negative = pipeline.predict_noise(**negative_kwargs)
+        actual = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=True,
+            true_cfg_scale=5.0,
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+
+    expected = negative + 5.0 * (positive - negative)
+    torch.testing.assert_close(actual, expected)
+
+
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
@@ -192,7 +192,6 @@ def test_constructor_weights_sources(boogu_pipeline):
        (DiffusionParallelConfig(tensor_parallel_size=2), "none", "Tensor parallelism"),
        (DiffusionParallelConfig(ulysses_degree=2), "none", "Sequence parallelism"),
        (DiffusionParallelConfig(ring_degree=2), "none", "Sequence parallelism"),
-        (DiffusionParallelConfig(cfg_parallel_size=2), "none", "CFG parallelism"),
        (
            DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=2),
            "none",
@@ -227,6 +226,26 @@ def test_constructor_rejects_unsupported_execution_modes(
    mock_dependencies["mllm_wrapper"].model.to.assert_not_called()


+@pytest.mark.parametrize("cfg_parallel_size", [2, 3])
+def test_constructor_accepts_cfg_parallel(mock_dependencies, cfg_parallel_size):
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        BooguImagePipeline,
+    )
+
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.float32,
+        parallel_config=DiffusionParallelConfig(cfg_parallel_size=cfg_parallel_size),
+    )
+
+    pipeline = BooguImagePipeline(od_config=od_config)
+
+    assert pipeline.od_config.parallel_config.cfg_parallel_size == cfg_parallel_size
+    assert hasattr(pipeline, "predict_noise_maybe_with_cfg")
+    assert hasattr(pipeline, "predict_noise_with_multi_branch_cfg")
+
+
 # ---------------------------------------------------------------------------
 # Prompt-encoding tests (deterministic fakes, no constructor)
 # ---------------------------------------------------------------------------
@@ -526,6 +545,14 @@ def test_forward_returns_diffusion_output():

 def test_forward_cfg_off_when_guidance_one():
    pipeline = _make_forward_pipeline()
+    cfg_calls = []
+    original_predict_noise_maybe_with_cfg = pipeline.predict_noise_maybe_with_cfg
+
+    def record_cfg_call(**kwargs):
+        cfg_calls.append(kwargs)
+        return original_predict_noise_maybe_with_cfg(**kwargs)
+
+    pipeline.predict_noise_maybe_with_cfg = record_cfg_call
    req = _make_request_batch("a cat", height=64, width=64, num_inference_steps=2, guidance_scale=1.0)
    # guidance_scale=1.0 is falsy-adjacent but explicitly disables CFG; the
    # request layer would set guidance_scale_provided, so emulate that here.
@@ -535,6 +562,30 @@ def test_forward_cfg_off_when_guidance_one():

    # Only the positive prompt is encoded when CFG is off.
    assert len(pipeline.processor.calls) == 1
+    assert len(cfg_calls) == 2
+    assert all(call["do_true_cfg"] is False for call in cfg_calls)
+    assert all(call["negative_kwargs"] is None for call in cfg_calls)
+
+
+def test_double_guidance_combine_matches_legacy_formula():
+    pipeline = _make_forward_pipeline()
+    positive_with_reference = torch.randn(2, 4, 8, 8)
+    negative_with_reference = torch.randn(2, 4, 8, 8)
+    uncond = torch.randn(2, 4, 8, 8)
+    text_scale = 5.0
+    image_scale = 2.0
+
+    actual = pipeline.combine_multi_branch_cfg_noise(
+        [positive_with_reference, negative_with_reference, uncond],
+        {"text": text_scale, "image": image_scale},
+    )
+    legacy = (
+        positive_with_reference
+        + (text_scale - 1.0) * (positive_with_reference - negative_with_reference)
+        + (image_scale - 1.0) * (negative_with_reference - uncond)
+    )
+
+    torch.testing.assert_close(actual, legacy, rtol=0, atol=1e-5)


 # ---------------------------------------------------------------------------
@@ -802,6 +853,39 @@ def test_build_ref_latents_expands_per_output_and_handles_none():
    assert ref_latents[2] is None and ref_latents[3] is None


+def test_vae_attention_context_is_noop_on_cpu():
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        BooguImagePipeline,
+    )
+
+    with BooguImagePipeline._vae_attention_context(torch.device("cpu")):
+        pass
+
+
+def test_vae_attention_context_prioritizes_efficient_cuda(monkeypatch):
+    from vllm_omni.diffusion.models.boogu_image import pipeline_boogu_image
+
+    recorded = {}
+
+    def fake_sdpa_kernel(backends, *, set_priority):
+        recorded["backends"] = backends
+        recorded["set_priority"] = set_priority
+        return pipeline_boogu_image.nullcontext()
+
+    monkeypatch.setattr(pipeline_boogu_image, "sdpa_kernel", fake_sdpa_kernel)
+
+    with pipeline_boogu_image.BooguImagePipeline._vae_attention_context(torch.device("cuda")):
+        pass
+
+    assert recorded == {
+        "backends": [
+            pipeline_boogu_image.SDPBackend.EFFICIENT_ATTENTION,
+            pipeline_boogu_image.SDPBackend.MATH,
+        ],
+        "set_priority": True,
+    }
+
+
 # ---------------------------------------------------------------------------
 # Editing / TI2I: forward CFG branch selection
 # ---------------------------------------------------------------------------
@@ -812,9 +896,11 @@ class _RecordingRefTransformer(_FakeTransformer):

    def __init__(self):
        self.calls = []
+        self.instruction_embed_calls = []

    def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
        self.calls.append(kwargs.get("ref_image_hidden_states"))
+        self.instruction_embed_calls.append(instruction_embeds)
        return torch.zeros_like(latents)


@@ -876,6 +962,9 @@ def test_forward_ti2i_text_only_two_predictions_with_ref():
    # Text-only ti2i: cond+ref and neg+ref -> 2 predictions/step, both carry ref.
    assert _count_per_step(pipeline.transformer.calls, num_steps) == 2
    assert all(ref is not None for ref in pipeline.transformer.calls)
+    for i in range(0, len(pipeline.transformer.instruction_embed_calls), 2):
+        positive, negative = pipeline.transformer.instruction_embed_calls[i : i + 2]
+        assert positive is not negative


 def test_forward_ti2i_double_guidance_three_predictions():
@@ -895,6 +984,11 @@ def test_forward_ti2i_double_guidance_three_predictions():
    per_step = [pipeline.transformer.calls[i : i + 3] for i in range(0, len(pipeline.transformer.calls), 3)]
    for step_calls in per_step:
        assert sum(ref is None for ref in step_calls) == 1
+    embed_calls = pipeline.transformer.instruction_embed_calls
+    for i in range(0, len(embed_calls), 3):
+        positive, negative_with_reference, uncond = embed_calls[i : i + 3]
+        assert positive is not negative_with_reference
+        assert negative_with_reference is uncond


 def test_forward_ti2i_image_only_two_predictions_drop_ref():
@@ -914,6 +1008,10 @@ def test_forward_ti2i_image_only_two_predictions_drop_ref():
    per_step = [pipeline.transformer.calls[i : i + 2] for i in range(0, len(pipeline.transformer.calls), 2)]
    for step_calls in per_step:
        assert sum(ref is None for ref in step_calls) == 1
+    embed_calls = pipeline.transformer.instruction_embed_calls
+    for i in range(0, len(embed_calls), 2):
+        positive_with_reference, positive_without_reference = embed_calls[i : i + 2]
+        assert positive_with_reference is positive_without_reference


 def test_forward_ti2i_no_guidance_single_prediction_with_ref():
@@ -6,7 +6,8 @@ L4 expansion coverage for Boogu-Image Edit and Edit-Turbo (image editing / TI2I)

 The nightly smoke module (``test_boogu_image_edit.py``) already covers single-image
 text-guided editing and the text+image double-guidance path. This expansion
-module adds the two edit paths smoke does not exercise:
+module adds the edit paths and CFG-parallel configurations smoke does not
+exercise:

 - ``test_image_only_guidance`` — ``guidance_scale=1.0`` (text CFG off) with
  ``guidance_scale_2=2.0`` (image CFG on): the image-only branch of upstream's
@@ -15,13 +16,19 @@ module adds the two edit paths smoke does not exercise:
 - ``test_align_res_output_size`` — a non-square reference with no requested
  ``height`` / ``width``: asserts the output resolution follows the reference
  (upstream ``align_res``, on by default for a single reference).
+- ``test_double_guidance_cfg_parallel`` — three-branch double guidance with
+  CFG sizes 2 and 3.
+- ``test_cfg_off_on_cfg_parallel_server`` — both scales at 1.0 on a size-2 CFG
+  server.

 Both Edit checkpoints share ``BooguImagePipeline`` with the Base checkpoint;
 the TI2I path activates when the request carries a reference image. Only a
 single reference image is supported. The regular Edit cases stay smoke-depth
 (``num_inference_steps=2``), while Edit-Turbo uses its tuned four-step
-configuration. CPU offload, Cache-DiT, and multi-GPU parallelism are not yet
-supported for these models and are intentionally omitted.
+configuration. Regular Edit CFG-parallel coverage includes two-card
+text-only/image-only/CFG-off paths, two-card round-robin double guidance, and
+full three-card double guidance. Edit-Turbo remains single-card-only here;
+CPU offload and Cache-DiT are intentionally omitted for both checkpoints.

 From ``tests/``::

@@ -47,16 +54,49 @@ EDIT_TURBO_REVISION = "hotfix-1k-20260708"
 EDIT_PROMPT = "Change the style to a colored pencil drawing."

 SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+PARALLEL_2_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+PARALLEL_3_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=3)


 def _get_diffusion_feature_cases(model: str):
-    """Return the single-card server configuration supported for Boogu-Image-Edit today."""
+    """Return single-card and two-card CFG configurations for two-branch paths."""
    return [
        pytest.param(
            OmniServerParams(model=model),
            id="default",
            marks=SINGLE_CARD_FEATURE_MARKS,
        ),
+        pytest.param(
+            OmniServerParams(model=model, server_args=["--cfg-parallel-size", "2"]),
+            id="cfg_parallel_2",
+            marks=PARALLEL_2_FEATURE_MARKS,
+        ),
+    ]
+
+
+def _get_double_guidance_cases(model: str):
+    """Return two-rank round-robin and full three-rank double-guidance rows."""
+    return [
+        pytest.param(
+            OmniServerParams(model=model, server_args=["--cfg-parallel-size", "2"]),
+            id="cfg_parallel_2_round_robin",
+            marks=PARALLEL_2_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(model=model, server_args=["--cfg-parallel-size", "3"]),
+            id="cfg_parallel_3_full",
+            marks=PARALLEL_3_FEATURE_MARKS,
+        ),
+    ]
+
+
+def _get_cfg_off_cases(model: str):
+    return [
+        pytest.param(
+            OmniServerParams(model=model, server_args=["--cfg-parallel-size", "2"]),
+            id="cfg_parallel_2_cfg_off",
+            marks=PARALLEL_2_FEATURE_MARKS,
+        )
    ]


@@ -106,6 +146,42 @@ def test_image_only_guidance(omni_server: OmniServer, openai_client: OpenAIClien
    openai_client.send_diffusion_request(request_config)


+@pytest.mark.parametrize("omni_server", _get_double_guidance_cases(MODEL), indirect=True)
+def test_double_guidance_cfg_parallel(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    """Double guidance exercises three branches on two and three CFG ranks."""
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "num_inference_steps": 2,
+            "guidance_scale": 5.0,
+            "guidance_scale_2": 2.0,
+            "seed": 42,
+        },
+    }
+    openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.parametrize("omni_server", _get_cfg_off_cases(MODEL), indirect=True)
+def test_cfg_off_on_cfg_parallel_server(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    """A CFG process group must not force negative embeddings or guidance."""
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "num_inference_steps": 2,
+            "guidance_scale": 1.0,
+            "guidance_scale_2": 1.0,
+            "seed": 42,
+        },
+    }
+    openai_client.send_diffusion_request(request_config)
+
+
@pytest.mark.parametrize(
    ("omni_server", "sampling_params"),
    _get_align_res_cases(),
@@ -7,14 +7,15 @@ L4 expansion coverage for ``Boogu/Boogu-Image-0.1-Base`` (text-to-image).
 The nightly smoke module (``test_boogu_image.py``) already covers the default
 512x512 single request and a concurrent batch. This expansion module adds the
 request-level paths that smoke does not exercise, run against every currently
-supported single-GPU server configuration:
+supported server configuration:

-Server rows (single card):
+Server rows:
 - ``default`` — no extra ``server_args``.
 - ``vae_slicing_tiling`` — ``--vae-use-slicing --vae-use-tiling`` (the recipe's
  OOM mitigation; applied generically in ``registry.initialize_model``).
+- ``cfg_parallel_2`` — two-card CFG branch parallelism.

-Cases (one per test, each parametrized over both server rows):
+Cases (one per test, each parametrized over all server rows):
 - ``test_non_square_resolution`` — 768x1024, exercises the patchify / RoPE path
  at a non-square aspect ratio.
 - ``test_cfg_off`` — ``guidance_scale=1.0``, exercises the no-CFG denoise branch
@@ -24,8 +25,7 @@ Cases (one per test, each parametrized over both server rows):
 Boogu-Image reads ``sp.guidance_scale`` (not ``true_cfg_scale``). These stay
 smoke-depth (``num_inference_steps=2``); numeric parity vs. Diffusers is covered
 by the local parity harness (support plan, step 14). CPU offload, Cache-DiT, and
-multi-GPU parallelism are not yet supported for this model and are intentionally
-omitted (add rows here when those features land).
+non-CFG multi-GPU parallelism remain unsupported.

 From ``tests/``::

@@ -49,10 +49,11 @@ T2I_PROMPT = "A mountain lake at sunset, photorealistic, cinematic lighting"
 NEGATIVE_PROMPT = ""

 SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+PARALLEL_2_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)


 def _get_diffusion_feature_cases(model: str):
-    """Return the single-card server configurations supported for Boogu-Image today."""
+    """Return the validated single-card and CFG-parallel configurations."""
    return [
        pytest.param(
            OmniServerParams(model=model),
@@ -67,6 +68,11 @@ def _get_diffusion_feature_cases(model: str):
            id="vae_slicing_tiling",
            marks=SINGLE_CARD_FEATURE_MARKS,
        ),
+        pytest.param(
+            OmniServerParams(model=model, server_args=["--cfg-parallel-size", "2"]),
+            id="cfg_parallel_2",
+            marks=PARALLEL_2_FEATURE_MARKS,
+        ),
    ]


@@ -16,11 +16,26 @@ from vllm_omni.diffusion.distributed.parallel_state import (
    get_cfg_group,
    get_classifier_free_guidance_rank,
    get_classifier_free_guidance_world_size,
+    is_cfg_group_initialized,
 )

 logger = init_logger(__name__)


+def _get_cfg_world_size_or_one() -> int:
+    """Return the CFG world size, defaulting only when no group exists.
+
+    Diffusion workers initialize the CFG process group even when its world
+    size is one.  Pipelines are also invoked directly by unit tests and some
+    offline integrations, though, where no distributed groups exist. Treat
+    that explicit case like a one-rank CFG group. Errors from an existing CFG
+    coordinator must propagate instead of silently selecting sequential CFG.
+    """
+    if not is_cfg_group_initialized():
+        return 1
+    return get_classifier_free_guidance_world_size()
+
+
 def _wrap(pred: torch.Tensor | tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
    """Normalize prediction to tuple form."""
    return pred if isinstance(pred, tuple) else (pred,)
@@ -106,7 +121,7 @@ class CFGParallelMixin(metaclass=ABCMeta):
        """
        if do_true_cfg:
            # Automatically detect CFG parallel configuration
-            cfg_parallel_ready = get_classifier_free_guidance_world_size() > 1
+            cfg_parallel_ready = _get_cfg_world_size_or_one() > 1

            if cfg_parallel_ready:
                cfg_group = get_cfg_group()
@@ -253,7 +268,7 @@ class CFGParallelMixin(metaclass=ABCMeta):
        """
        if do_true_cfg:
            n_branches = len(branches_kwargs)
-            cfg_world_size = get_classifier_free_guidance_world_size()
+            cfg_world_size = _get_cfg_world_size_or_one()
            cfg_parallel_ready = cfg_world_size > 1

            if cfg_parallel_ready:
@@ -474,7 +489,7 @@ class CFGParallelMixin(metaclass=ABCMeta):
        Returns:
            True if CFG parallel configuration is valid, False otherwise
        """
-        if get_classifier_free_guidance_world_size() == 1:
+        if _get_cfg_world_size_or_one() == 1:
            return True

        if true_cfg_scale <= 1:
@@ -322,6 +322,11 @@ def is_pipeline_last_stage():


 # CFG
+def is_cfg_group_initialized() -> bool:
+    """Return whether the classifier-free-guidance group has been created."""
+    return _CFG is not None
+
+
 def get_cfg_group() -> GroupCoordinator:
    assert _CFG is not None, "classifier_free_guidance parallel group is not initialized"
    return _CFG
@@ -11,17 +11,17 @@ Ported from the upstream ``boogu`` package
  loaded from the checkpoint subfolders; transformer weights arrive later via
  ``weights_sources`` + ``load_weights``).
 - Upstream ``encode_instruction`` is exposed as ``encode_prompt`` (the
-  vLLM-Omni convention, also hooked by the prompt-embed cache); only the
-  text-to-image path is ported. Instruction rewriting, prompt tuning,
-  reference-image (TI2I/I2I) encoding, double guidance, and vision-token
-  stripping are not ported.
- ``forward()`` (denoise loop + VAE decode) is intentionally not implemented
-  yet; it lands together with the post-process registration.
+  vLLM-Omni convention, also hooked by the prompt-embed cache).
+- Text-to-image and single-reference TI2I inference share one native pipeline;
+  CFG branches use vLLM-Omni's shared two-branch/N-branch parallel helpers.
+- Instruction rewriting, prompt tuning, and vision-token stripping are not
+  ported.
 """

 import json
 import os
 from collections.abc import Iterable
+from contextlib import nullcontext
 from typing import ClassVar, cast

 import PIL.Image
@@ -31,11 +31,13 @@ from diffusers.image_processor import VaeImageProcessor
 from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
+from torch.nn.attention import SDPBackend, sdpa_kernel
 from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
 from vllm.logger import init_logger
 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.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
@@ -189,7 +191,7 @@ SYSTEM_PROMPT_4_T2I_UNIFIED = (
 )


-class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery, SupportImageInput):
+class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, SupportsComponentDiscovery, SupportImageInput):
    """Boogu-Image text-to-image and image-editing (TI2I) pipeline.

    Native vLLM-Omni implementation. A request with a reference image (edit /
@@ -299,8 +301,6 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
            raise NotImplementedError("Tensor parallelism is not supported by BooguImagePipeline.")
        if (parallel_config.sequence_parallel_size or 1) > 1:
            raise NotImplementedError("Sequence parallelism is not supported by BooguImagePipeline.")
-        if parallel_config.cfg_parallel_size > 1:
-            raise NotImplementedError("CFG parallelism is not supported by BooguImagePipeline.")
        if parallel_config.use_hsdp:
            raise NotImplementedError("HSDP is not supported by BooguImagePipeline.")
        if self.od_config.cache_backend not in (None, "", "none"):
@@ -562,6 +562,79 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
            ref_image_hidden_states=ref_image_hidden_states,
        )

+    def predict_noise(self, **kwargs) -> torch.Tensor:
+        """Run one Boogu CFG branch through the native transformer."""
+        return self.predict(**kwargs)
+
+    def combine_cfg_noise(
+        self,
+        positive_noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
+        negative_noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
+        true_cfg_scale: float,
+        cfg_normalize: bool = False,
+        kwargs: dict | None = None,
+    ) -> torch.Tensor:
+        """Preserve Boogu's sequential two-branch CFG operation order."""
+        positive_items = positive_noise_pred if isinstance(positive_noise_pred, tuple) else (positive_noise_pred,)
+        negative_items = negative_noise_pred if isinstance(negative_noise_pred, tuple) else (negative_noise_pred,)
+        if len(positive_items) != 1 or len(negative_items) != 1:
+            raise ValueError("Boogu CFG expects exactly one prediction tensor per branch.")
+
+        positive = positive_items[0]
+        negative = negative_items[0]
+        combined = positive + (true_cfg_scale - 1) * (positive - negative)
+        if cfg_normalize:
+            combined = self.cfg_normalize_function(positive, combined)
+        return combined
+
+    def combine_multi_branch_cfg_noise(
+        self,
+        predictions: list[torch.Tensor],
+        true_cfg_scale: float | dict[str, float],
+        cfg_normalize: bool = False,
+    ) -> torch.Tensor:
+        """Combine Boogu CFG branches using the original operation order.
+
+        Although the usual two- and three-branch CFG formulas can be rewritten
+        algebraically, changing their floating-point operation order causes
+        small per-step differences that accumulate across the denoise loop.
+        Keep the sequential Boogu implementation's order so parallel branch
+        combination does not introduce an additional source of numeric drift.
+
+        The two-branch formula is::
+
+            positive + (scale - 1) * (positive - negative)
+
+        The three-branch order is ``[positive_with_reference,
+        negative_with_reference, negative_without_reference]`` and combines as::
+
+            positive_with_reference
+            + (text_scale - 1) * (positive_with_reference - negative_with_reference)
+            + (image_scale - 1) * (negative_with_reference - uncond)
+        """
+        if len(predictions) == 2:
+            if not isinstance(true_cfg_scale, float):
+                raise TypeError("Boogu two-branch CFG requires a scalar guidance scale.")
+            positive, negative = predictions
+            combined = positive + (true_cfg_scale - 1) * (positive - negative)
+            if cfg_normalize:
+                combined = self.cfg_normalize_function(positive, combined)
+            return combined
+        if len(predictions) != 3:
+            return super().combine_multi_branch_cfg_noise(predictions, true_cfg_scale, cfg_normalize)
+        if not isinstance(true_cfg_scale, dict):
+            raise TypeError("Boogu three-branch CFG requires text and image guidance scales.")
+
+        positive_with_reference, negative_with_reference, uncond = predictions
+        combined = (
+            positive_with_reference
+            + (true_cfg_scale["text"] - 1) * (positive_with_reference - negative_with_reference)
+            + (true_cfg_scale["image"] - 1) * (negative_with_reference - uncond)
+        )
+        if cfg_normalize:
+            combined = self.cfg_normalize_function(positive_with_reference, combined)
+        return combined
+
    def _encode_vae_image(self, img: torch.Tensor, generator=None) -> torch.Tensor:
        """Encode an image tensor into the VAE latent space (upstream ``encode_vae``).

@@ -569,13 +642,35 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
        threads the request generator through so a fixed seed gives a
        reproducible reference latent.
        """
-        z0 = self.vae.encode(img.to(dtype=self.vae.dtype)).latent_dist.sample(generator=generator)
+        with self._vae_attention_context(img.device):
+            z0 = self.vae.encode(img.to(dtype=self.vae.dtype)).latent_dist.sample(generator=generator)
        if self.vae.config.shift_factor is not None:
            z0 = z0 - self.vae.config.shift_factor
        if self.vae.config.scaling_factor is not None:
            z0 = z0 * self.vae.config.scaling_factor
        return z0.to(dtype=self.vae.dtype)

+    @staticmethod
+    def _vae_attention_context(device: torch.device):
+        """Pin the VAE SDPA backend so worker topology cannot change results.
+
+        Boogu's float32 VAE attention can otherwise choose a different SDPA
+        backend in the in-process CFG=1 executor and the multi-process CFG>1
+        executor. The resulting small reference-latent differences accumulate
+        over the denoising loop. Prefer the default fast CUDA backend used by
+        Boogu and keep the math implementation as a CUDA fallback.
+
+        This pin is intentionally CUDA-only. CPU and vendor-specific accelerator
+        backends keep their platform defaults, so topology-independent VAE
+        parity is not guaranteed for those devices by this context.
+        """
+        if device.type != "cuda":
+            return nullcontext()
+        return sdpa_kernel(
+            [SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH],
+            set_priority=True,
+        )
+
    def _build_ref_latents(
        self,
        preprocessed_images: list[torch.Tensor | None],
@@ -722,71 +817,107 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
        timesteps = self.scheduler.timesteps
        num_timesteps = len(timesteps)

-        # 5. Denoise loop with (sequential) classifier-free guidance. Reproduces
-        # the branch priority of upstream ``processing`` (double > text-only >
-        # image-only > t2i text). Reference latents are kept in the conditional
-        # (and, for text-only ti2i, the unconditional) predictions.
+        # 5. Denoise loop with shared sequential/parallel CFG execution.
+        # Reproduces the branch priority of upstream ``processing`` (double >
+        # text-only > image-only > t2i text). Reference latents stay attached
+        # to the same branches as in the original sequential implementation.
        with self.progress_bar(total=num_timesteps) as progress_bar:
            for i, t in enumerate(timesteps):
                in_cfg_range = cfg_range[0] <= i / num_timesteps <= cfg_range[1]
                text_gs = text_guidance_scale if in_cfg_range else 1.0
                image_gs = image_guidance_scale if in_cfg_range else 1.0

-                model_pred = self.predict(
-                    t, latents, instruction_embeds, freqs_real, instruction_attention_mask, ref_latents
+                positive_kwargs = dict(
+                    t=t,
+                    latents=latents,
+                    instruction_embeds=instruction_embeds,
+                    freqs_real=freqs_real,
+                    instruction_attention_mask=instruction_attention_mask,
+                    ref_image_hidden_states=ref_latents,
                )

                if task_type == "ti2i" and text_gs > 1.0 and image_gs > 1.0:
                    # Double guidance: 3 predictions (cond+ref, neg+ref, neg+no-ref).
-                    model_pred_drop_text = self.predict(
-                        t,
-                        latents,
-                        negative_instruction_embeds,
-                        freqs_real,
-                        negative_instruction_attention_mask,
-                        ref_latents,
+                    negative_with_reference_kwargs = dict(
+                        t=t,
+                        latents=latents,
+                        instruction_embeds=negative_instruction_embeds,
+                        freqs_real=freqs_real,
+                        instruction_attention_mask=negative_instruction_attention_mask,
+                        ref_image_hidden_states=ref_latents,
                    )
-                    model_pred_drop_all = self.predict(
-                        t,
-                        latents,
-                        negative_instruction_embeds,
-                        freqs_real,
-                        negative_instruction_attention_mask,
-                        None,
+                    uncond_kwargs = dict(
+                        t=t,
+                        latents=latents,
+                        instruction_embeds=negative_instruction_embeds,
+                        freqs_real=freqs_real,
+                        instruction_attention_mask=negative_instruction_attention_mask,
+                        ref_image_hidden_states=None,
+                    )
+                    model_pred = self.predict_noise_with_multi_branch_cfg(
+                        do_true_cfg=True,
+                        true_cfg_scale={"text": text_gs, "image": image_gs},
+                        branches_kwargs=[positive_kwargs, negative_with_reference_kwargs, uncond_kwargs],
+                        cfg_normalize=False,
                    )
-                    delta_text = model_pred - model_pred_drop_text
-                    delta_image = model_pred_drop_text - model_pred_drop_all
-                    model_pred = model_pred + (text_gs - 1) * delta_text + (image_gs - 1) * delta_image
                elif task_type == "ti2i" and text_gs > 1.0:
                    # Text-only ti2i guidance: reference kept in the uncond pred.
-                    model_pred_drop_text = self.predict(
-                        t,
-                        latents,
-                        negative_instruction_embeds,
-                        freqs_real,
-                        negative_instruction_attention_mask,
-                        ref_latents,
+                    negative_kwargs = dict(
+                        t=t,
+                        latents=latents,
+                        instruction_embeds=negative_instruction_embeds,
+                        freqs_real=freqs_real,
+                        instruction_attention_mask=negative_instruction_attention_mask,
+                        ref_image_hidden_states=ref_latents,
+                    )
+                    model_pred = self.predict_noise_maybe_with_cfg(
+                        do_true_cfg=True,
+                        true_cfg_scale=text_gs,
+                        positive_kwargs=positive_kwargs,
+                        negative_kwargs=negative_kwargs,
+                        cfg_normalize=False,
                    )
-                    model_pred = model_pred + (text_gs - 1) * (model_pred - model_pred_drop_text)
                elif task_type == "ti2i" and image_gs > 1.0:
                    # Image-only ti2i guidance: drop the reference in the uncond pred.
-                    model_pred_drop_image = self.predict(
-                        t, latents, instruction_embeds, freqs_real, instruction_attention_mask, None
+                    negative_kwargs = dict(positive_kwargs, ref_image_hidden_states=None)
+                    model_pred = self.predict_noise_maybe_with_cfg(
+                        do_true_cfg=True,
+                        true_cfg_scale=image_gs,
+                        positive_kwargs=positive_kwargs,
+                        negative_kwargs=negative_kwargs,
+                        cfg_normalize=False,
                    )
-                    model_pred = model_pred + (image_gs - 1) * (model_pred - model_pred_drop_image)
                elif text_gs > 1.0:
                    # Text-to-image classifier-free guidance.
-                    model_pred_drop_all = self.predict(
-                        t,
-                        latents,
-                        negative_instruction_embeds,
-                        freqs_real,
-                        negative_instruction_attention_mask,
-                        None,
+                    negative_kwargs = dict(
+                        t=t,
+                        latents=latents,
+                        instruction_embeds=negative_instruction_embeds,
+                        freqs_real=freqs_real,
+                        instruction_attention_mask=negative_instruction_attention_mask,
+                        ref_image_hidden_states=None,
+                    )
+                    model_pred = self.predict_noise_maybe_with_cfg(
+                        do_true_cfg=True,
+                        true_cfg_scale=text_gs,
+                        positive_kwargs=positive_kwargs,
+                        negative_kwargs=negative_kwargs,
+                        cfg_normalize=False,
+                    )
+                else:
+                    # CFG-off requests remain valid even if the server owns a
+                    # CFG process group: every rank evaluates only the positive
+                    # branch and no negative embeddings are required.
+                    model_pred = self.predict_noise_maybe_with_cfg(
+                        do_true_cfg=False,
+                        true_cfg_scale=1.0,
+                        positive_kwargs=positive_kwargs,
+                        negative_kwargs=None,
+                        cfg_normalize=False,
                    )
-                    model_pred = model_pred + (text_gs - 1) * (model_pred - model_pred_drop_all)

-                latents = self.scheduler.step(model_pred, t, latents, return_dict=False)[0]
+                do_true_cfg = text_gs > 1.0 or (task_type == "ti2i" and image_gs > 1.0)
+                latents = self.scheduler_step_maybe_with_cfg(model_pred, t, latents, do_true_cfg=do_true_cfg)
                latents = latents.to(dtype=instruction_embeds.dtype)
                progress_bar.update()

@@ -799,7 +930,8 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery
                latents = latents / self.vae.config.scaling_factor
            if self.vae.config.shift_factor is not None:
                latents = latents + self.vae.config.shift_factor
-            image = self.vae.decode(latents, return_dict=False)[0]
+            with self._vae_attention_context(latents.device):
+                image = self.vae.decode(latents, return_dict=False)[0]
            if (ori_height, ori_width) != (height, width):
                image = F.interpolate(image, size=(ori_height, ori_width), mode="bilinear")