[Feat] Support CFG-Parallel For Sensenova-U1-8B-MoT (#4188)

Signed-off-by: sphinxkkkbc <binchengkang8@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: sphinxkkkbc <binchengkang8@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>

tests/diffusion/distributed/test_cfg_parallel.py

@@ -208,6 +208,7 @@ def _test_cfg_parallel_worker(
             positive_kwargs=positive_kwargs,
             negative_kwargs=negative_kwargs,
             cfg_normalize=test_config["cfg_normalize"],
+            kwargs=test_config["kwargs"],
         )
 
     # CFG parallel returns the combined prediction on every rank.
@@ -284,6 +285,7 @@ def _test_cfg_sequential_worker(
             positive_kwargs=positive_kwargs,
             negative_kwargs=negative_kwargs,
             cfg_normalize=test_config["cfg_normalize"],
+            kwargs=test_config["kwargs"],
         )
 
     # Sequential CFG always returns output
@@ -297,7 +299,10 @@ def _test_cfg_sequential_worker(
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
 @pytest.mark.parametrize("batch_size", [2])
 @pytest.mark.parametrize("cfg_normalize", [False, True])
-def test_predict_noise_maybe_with_cfg(cfg_parallel_size: int, dtype: torch.dtype, batch_size: int, cfg_normalize: bool):
+@pytest.mark.parametrize("kwargs", [None, {"step_i": 1}])
+def test_predict_noise_maybe_with_cfg(
+    cfg_parallel_size: int, dtype: torch.dtype, batch_size: int, cfg_normalize: bool, kwargs: dict | None
+):
     """
     Test that predict_noise_maybe_with_cfg produces identical results
     with and without CFG parallel.
@@ -322,6 +327,7 @@ def test_predict_noise_maybe_with_cfg(cfg_parallel_size: int, dtype: torch.dtype
         "cfg_normalize": cfg_normalize,
         "model_seed": 42,  # Fixed seed for model initialization
         "input_seed": 123,  # Fixed seed for input generation
+        "kwargs": kwargs,  # Additional kwargs to test passing through CFG parallel
     }
 
     mp_context = torch.multiprocessing.get_context("spawn")

tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import SenseNovaU1Pipeline
+
+
+def _pipeline_without_init() -> SenseNovaU1Pipeline:
+    return object.__new__(SenseNovaU1Pipeline)
+
+
+@pytest.mark.parametrize("kwargs", [None, {"step_i": 0}])
+def test_combine_cfg_noise_requires_is_it2i(kwargs):
+    pipe = _pipeline_without_init()
+    out_cond = (torch.ones(1, 2, 3),)
+    out_uncond = (torch.zeros(1, 2, 3),)
+
+    with pytest.raises(ValueError, match="is_it2i"):
+        pipe.combine_cfg_noise(
+            out_cond,
+            out_uncond,
+            cfg_scale=4.0,
+            cfg_norm="cfg_zero_star",
+            kwargs=kwargs,
+        )
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [{"is_it2i": False}, {"is_it2i": False, "step_i": None}],
+)
+def test_cfg_zero_star_requires_step_i(kwargs):
+    pipe = _pipeline_without_init()
+    out_cond = (torch.ones(1, 2, 3),)
+    out_uncond = (torch.zeros(1, 2, 3),)
+    with pytest.raises(ValueError, match="step_i"):
+        pipe.combine_cfg_noise(
+            out_cond,
+            out_uncond,
+            cfg_scale=4.0,
+            cfg_norm="cfg_zero_star",
+            kwargs=kwargs,
+        )
+
+
+def test_cfg_zero_star_accepts_step_i():
+    pipe = _pipeline_without_init()
+    out_cond = (torch.ones(1, 2, 3),)
+    out_uncond = (torch.zeros(1, 2, 3),)
+    result = pipe.combine_cfg_noise(
+        out_cond,
+        out_uncond,
+        cfg_scale=4.0,
+        cfg_norm="cfg_zero_star",
+        kwargs={"is_it2i": False, "step_i": 0},
+    )
+
+    assert result.shape == out_cond[0].shape
+    assert torch.equal(result, torch.zeros_like(out_cond[0]))
+    assert torch.isfinite(result).all()