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

Signed-off-by: nagisa <1434936049@qq.com>
Signed-off-by: nagisa-kun <1434936049@qq.com>

tests/diffusion/distributed/test_cfg_parallel.py

@@ -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

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -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():