[Feat] [CI] [bugfix] Generalize LTX CFG parallelism to the complete guidance plan (#5547)

Signed-off-by: mglyn <1203789601@qq.com>

recipes/LTX/LTX-2.md

@@ -165,6 +165,40 @@ Pass these fields in online `extra_params` or offline `--extra-body`:
 --extra-body '{"video_cfg_scale":3.0,"audio_cfg_scale":7.0}'
 ```
 
+### Guidance Parallelism
+
+For LTX, `cfg_parallel_size` is the number of ranks used to execute the
+complete guidance plan in parallel. Despite the legacy `cfg` name, it covers
+text CFG, STG, and cross-modality guidance passes; guidance rescaling is
+applied after all pass predictions are gathered.
+
+The default one-stage recipe has four Transformer passes per denoise step:
+`cond`, `uncond`, `ptb` (STG), and `mod` (cross-modality). The useful balanced
+configurations are therefore:
+
+| `--cfg-parallel-size` | Passes per rank | Guidance-slot utilization | Notes |
+|---:|---:|---:|---|
+| `1` | 4 | 100% | Single-rank fused guidance batch |
+| `2` | 2 | 100% | Recommended two-rank configuration |
+| `4` | 1 | 100% | One guidance pass per rank |
+
+Other positive sizes are accepted. When the pass count is not divisible by
+`cfg_parallel_size`, ranks are padded to an equal number of execution slots
+and LTX emits a warning with the expected utilization. For example, four
+passes on three ranks use six slots and have 66.7% guidance-slot utilization.
+
+Start an LTX-2.3 server with two-way guidance parallelism:
+
+```bash
+vllm serve diffusers/LTX-2.3-Diffusers --omni \
+  --cfg-parallel-size 2 --stage-init-timeout 600
+```
+
+The total device count is the product of `cfg_parallel_size` and the other
+configured parallel dimensions. Distilled pipelines use positive-only
+guidance, so `cfg_parallel_size > 1` does not provide guidance speedup for
+them.
+
 ### Python API
 
 Normal per-request values belong in `OmniDiffusionSamplingParams`. Pass the
@@ -273,8 +307,8 @@ bundled offline CLI do not currently expose `sigmas`.
 - Fused one-stage requests must resolve to identical LTX guidance and sigma
   schedules. Keep concurrent requests guidance-homogeneous or use
   `--max-num-seqs 1`.
-- `--cfg-parallel-size 2` supports only the CFG-only plan: disable STG and
-  modality guidance and set both rescale fields to `0.0`.
+- `--cfg-parallel-size` shards the complete LTX guidance plan, including STG,
+  modality guidance, and rescale-compatible prediction gathering.
 
 ## Operational Notes
 

tests/config/test_omni_config.py

@@ -577,9 +577,11 @@ def test_from_pipeline_config_derives_sequence_parallel_size_from_allgather_degr
     assert stage.parallel_config.world_size == 2
 
 
-def test_diffusion_parallel_config_rejects_cfg_parallel_size_outside_current_bound():
-    with pytest.raises(ValidationError):
-        OmniStageDiffusionParallelConfig(cfg_parallel_size=4)
+def test_diffusion_parallel_config_accepts_four_way_guidance_parallelism():
+    cfg = OmniStageDiffusionParallelConfig(cfg_parallel_size=4)
+
+    assert cfg.cfg_parallel_size == 4
+    assert cfg.world_size == 4
 
 
 def test_diffusion_parallel_config_rejects_allgather_with_ulysses_or_ring():

tests/diffusion/models/ltx2/test_ltx2_conditioning.py

@@ -110,7 +110,7 @@ class TestLTXImageToVideoForwardStages:
 
         kwargs = pipe._denoise_timestep_kwargs(
             ts,
-            SimpleNamespace(cfg_parallel_ready=False),
+            SimpleNamespace(guidance_parallel_ready=False),
             denoise_ctx,
             video_token_count=2,
             audio_token_count=1,

tests/diffusion/models/ltx2/test_ltx2_parallel.py

@@ -3,6 +3,7 @@
 
 """Unit tests for LTX guidance and forward-parallel behavior."""
 
+from contextlib import nullcontext
 from types import SimpleNamespace
 
 import pytest
@@ -50,7 +51,8 @@ class TestCFGParallelHelpers:
         assert torch.allclose(video_combined, expected_video)
         assert torch.allclose(audio_combined, expected_audio)
 
-    def test_cfg_parallel_rejects_non_cfg_guidance(self):
+    @pytest.mark.parametrize("guidance_world_size", [1, 2, 3, 4])
+    def test_guidance_parallel_accepts_official_guidance(self, guidance_world_size):
         from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
             LTX_GUIDANCE_EXECUTOR,
             LTXGuidancePlan,
@@ -59,26 +61,180 @@ class TestCFGParallelHelpers:
 
         plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance)
 
-        with pytest.raises(ValueError, match="CFG-only guidance"):
-            LTX_GUIDANCE_EXECUTOR.validate_cfg_world_size(plan, 2)
+        LTX_GUIDANCE_EXECUTOR.validate_guidance_world_size(plan, guidance_world_size)
 
-    def test_cfg_parallel_dummy_warms_supported_cfg_only_plan(self, monkeypatch):
+    def test_guidance_parallel_dummy_warms_full_guidance_plan(self, monkeypatch):
         from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
         from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
-        monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 2)
         pipe = object.__new__(LTX2Pipeline)
         req = SimpleNamespace(is_dummy_run=lambda: True)
         request_inputs = SimpleNamespace(guidance=LTX23_ONE_STAGE_RECIPE.request_guidance)
 
-        cfg_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None)
+        guidance_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None)
 
-        assert cfg_parallel_ready is True
-        assert pipe._guidance_plan.names == ("cond", "uncond")
+        assert guidance_parallel_ready is True
+        assert pipe._guidance_plan.names == ("cond", "uncond", "ptb", "mod")
         assert pipe._guidance_plan.spec.video.cfg_scale == 3.0
         assert pipe._guidance_plan.spec.audio.cfg_scale == 7.0
 
+    def test_guidance_parallel_warns_for_imbalanced_passes(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2 import ltx2_guidance
+
+        spec = ltx2_guidance.LTXGuidanceSpec(
+            video=ltx2_guidance.LTXModalityGuidance(
+                cfg_scale=3.0,
+                stg_scale=1.0,
+                stg_blocks=(1,),
+            )
+        )
+        plan = ltx2_guidance.LTXGuidancePlan.build(spec)
+        captured = []
+        monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_rank", lambda: 0)
+        monkeypatch.setattr(ltx2_guidance.logger, "warning_once", lambda *args: captured.append(args))
+
+        ltx2_guidance.LTX_GUIDANCE_EXECUTOR.warn_if_imbalanced(plan, 4, "hq_generate")
+
+        assert plan.names == ("cond", "uncond", "ptb")
+        assert len(captured) == 1
+        assert captured[0][1:] == ("hq_generate", 3, 4, 1, 75.0)
+
+    @pytest.mark.parametrize("guidance_rank", [0, 1])
+    def test_default_guidance_is_sharded_and_combined(self, monkeypatch, guidance_rank):
+        from vllm_omni.diffusion.models.ltx2 import ltx2_guidance
+        from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext
+        from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
+        from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE
+
+        executor = ltx2_guidance.LTX_GUIDANCE_EXECUTOR
+        plan = ltx2_guidance.LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance)
+        video_predictions = {
+            "cond": torch.tensor([[[0.2, -0.3]]]),
+            "uncond": torch.tensor([[[-0.4, 0.1]]]),
+            "ptb": torch.tensor([[[0.5, -0.2]]]),
+            "mod": torch.tensor([[[-0.1, 0.6]]]),
+        }
+        audio_predictions = {
+            "cond": torch.tensor([[[0.7, -0.2]]]),
+            "uncond": torch.tensor([[[0.1, 0.4]]]),
+            "ptb": torch.tensor([[[-0.3, 0.2]]]),
+            "mod": torch.tensor([[[0.4, -0.5]]]),
+        }
+        state = LTXAVState(
+            video=torch.tensor([[[1.0, -2.0]]]),
+            audio=torch.tensor([[[0.5, 3.0]]]),
+        )
+        denoise_ctx = LTXDenoiseContext(
+            latents=state.video,
+            audio_latents=state.audio,
+            video_coords=torch.zeros(1, 1, 3),
+            audio_coords=torch.zeros(1, 1, 1),
+        )
+        executor.prepare_denoise_context(plan, True, 2, denoise_ctx)
+
+        class FakeGuidanceGroup:
+            def __init__(self):
+                self.call_index = 0
+
+            def all_gather(self, tensor, separate_tensors=True):
+                assert separate_tensors
+                gathered = (
+                    [video_predictions["cond"], video_predictions["uncond"]],
+                    [video_predictions["ptb"], video_predictions["mod"]],
+                    [audio_predictions["cond"], audio_predictions["uncond"]],
+                    [audio_predictions["ptb"], audio_predictions["mod"]],
+                )[self.call_index]
+                torch.testing.assert_close(tensor, gathered[guidance_rank])
+                self.call_index += 1
+                return gathered
+
+        group = FakeGuidanceGroup()
+        monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_rank", lambda: guidance_rank)
+        monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_group", lambda: group)
+
+        local_names = ("cond", "ptb") if guidance_rank == 0 else ("uncond", "mod")
+        captured = {}
+
+        class FakeTransformer:
+            def __call__(self, **kwargs):
+                captured.update(kwargs)
+                return (
+                    torch.cat([video_predictions[name] for name in local_names]),
+                    torch.cat([audio_predictions[name] for name in local_names]),
+                )
+
+        video_sigma = torch.tensor(0.25)
+        audio_sigma = torch.tensor(0.5)
+
+        class FakePipeline:
+            scheduler = SimpleNamespace(sigmas=torch.stack([video_sigma]))
+            transformer = FakeTransformer()
+
+            @staticmethod
+            def _transformer_cache_context(name):
+                assert name == "guided"
+                return nullcontext()
+
+            @staticmethod
+            def _build_transformer_kwargs(forward_ctx, denoise_ctx, **kwargs):
+                del forward_ctx, denoise_ctx
+                return kwargs
+
+            @staticmethod
+            def _video_guidance_model_sigma(sigma, denoise_ctx):
+                del denoise_ctx
+                return sigma
+
+        positive_context = torch.ones(1, 1, 1)
+        negative_context = torch.zeros(1, 1, 1)
+        forward_ctx = SimpleNamespace(
+            prompt_context=SimpleNamespace(
+                positive_connector_prompt_embeds=positive_context,
+                positive_connector_audio_prompt_embeds=positive_context,
+                negative_connector_prompt_embeds=negative_context,
+                negative_connector_audio_prompt_embeds=negative_context,
+            ),
+            attention_kwargs=None,
+            audio_scheduler=SimpleNamespace(sigmas=torch.stack([audio_sigma])),
+        )
+
+        actual_video, actual_audio = executor.predict_parallel_guidance(
+            FakePipeline(),
+            plan,
+            index=0,
+            timestep=torch.tensor(1.0),
+            state=state,
+            forward_ctx=forward_ctx,
+            denoise_ctx=denoise_ctx,
+        )
+
+        expected_video = executor._guide_modality(
+            state.video,
+            video_predictions,
+            video_sigma,
+            plan.spec.video,
+        )
+        expected_audio = executor._guide_modality(
+            state.audio,
+            audio_predictions,
+            audio_sigma,
+            plan.spec.audio,
+        )
+        torch.testing.assert_close(actual_video, expected_video)
+        torch.testing.assert_close(actual_audio, expected_audio)
+        assert group.call_index == 4
+        assert captured["hidden_states"].shape[0] == 2
+        perturbations = captured["attention_kwargs"]["ltx_perturbation_kwargs"]
+        if guidance_rank == 0:
+            assert "video_self_attention_mask" in perturbations
+            assert "a2v_cross_attention_mask" not in perturbations
+        else:
+            assert "video_self_attention_mask" not in perturbations
+            assert "a2v_cross_attention_mask" in perturbations
+
 
 class TestCFGParallelForwardPath:
     """Test the LTX-2.3 CFG-parallel denoising path without loading model weights."""
@@ -96,7 +252,7 @@ class TestCFGParallelForwardPath:
         pipe.tokenizer_max_length = 4
         pipe.vae_spatial_compression_ratio = 32
         pipe.vae_temporal_compression_ratio = 8
-        monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 1)
+        monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 1)
 
         class StopAtEncodePromptError(Exception):
             pass
@@ -238,10 +394,10 @@ class TestCFGParallelForwardPath:
                     return [audio_pos, audio_neg]
                 raise AssertionError(f"Unexpected gathered tensor: {tensor}")
 
-        monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 2)
-        monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2)
-        monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda: cfg_rank)
-        monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup())
+        monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_rank", lambda: cfg_rank)
+        monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_group", lambda: FakeCfgGroup())
 
         def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None):
             scheduler.sigmas = torch.tensor([0.25, 0.25], device=device)

tests/entrypoints/test_serve.py

@@ -46,6 +46,16 @@ def test_serve_parser_accepts_strategy_config() -> None:
     assert args.get_explicit_kwargs_dict()["strategy_config"] == "/tmp/strategy.yaml"
 
 
+def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(["serve", "fake-model", "--omni", "--cfg-parallel-size", "4"])
+
+    assert args.cfg_parallel_size == 4
+
+
 def _make_headless_args(**kwargs) -> TrackingNamespace:
     defaults = {
         "model": "fake-model",

vllm_omni/config/omni_config.py

@@ -368,7 +368,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
     ring_degree: int = Field(default=1, ge=1)
     allgather_degree: int = Field(default=1, ge=1)
     ulysses_mode: str = "strict"
-    cfg_parallel_size: int = Field(default=1, ge=1, le=3)
+    cfg_parallel_size: int = Field(default=1, ge=1)
     vae_patch_parallel_size: int = Field(default=1, ge=1)
     vae_parallel_mode: str = "tile"
     use_hsdp: bool = False

vllm_omni/diffusion/data.py

@@ -185,7 +185,7 @@ class DiffusionParallelConfig:
     """
 
     cfg_parallel_size: int = 1
-    """Number of Classifier Free Guidance (CFG) parallel groups."""
+    """Number of ranks used to execute guidance passes in parallel."""
 
     vae_patch_parallel_size: int = 1
     """Number of ranks used for VAE patch/tile parallelism (decode/encode)."""
@@ -233,9 +233,6 @@ class DiffusionParallelConfig:
         assert self.ring_degree > 0, "Ring degree must be > 0"
         assert self.allgather_degree > 0, "AllGather degree must be > 0"
         assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0"
-        assert self.cfg_parallel_size in [1, 2, 3], (
-            f"CFG parallel size must be 1, 2, or 3, but got {self.cfg_parallel_size}"
-        )
         assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0"
         assert self.vae_parallel_mode in {"tile", "spatial_shard_height", "spatial_shard_width"}, (
             "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "

vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py

@@ -705,12 +705,12 @@ class LTXI2VConditioningMixin:
             latents=request_inputs.latents,
         )
 
-    def _prepare_denoise_context_for_cfg(
+    def _prepare_denoise_context_for_guidance(
         self,
         forward_ctx: LTXForwardContext,
         denoise_ctx: LTXDenoiseContext,
     ) -> LTXDenoiseContext:
-        denoise_ctx = super()._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx)
+        denoise_ctx = super()._prepare_denoise_context_for_guidance(forward_ctx, denoise_ctx)
         if denoise_ctx.conditioning_mask is None:
             if self.unified_text_image_entry:
                 return denoise_ctx
@@ -755,9 +755,7 @@ class LTXI2VConditioningMixin:
             video_token_count=video_token_count,
             audio_token_count=audio_token_count,
         )
-        conditioning_mask = (
-            denoise_ctx.conditioning_mask if forward_ctx.cfg_parallel_ready else denoise_ctx.conditioning_mask_for_model
-        )
+        conditioning_mask = denoise_ctx.conditioning_mask_for_model
         if conditioning_mask is None:
             raise ValueError("LTX I2V denoising requires a conditioning mask.")
         kwargs.update(

vllm_omni/diffusion/models/ltx2/ltx2_denoise.py

@@ -33,7 +33,7 @@ class LTXForwardContext:
     request_inputs: LTXRequestInputs
     prompt_context: LTXPromptContext
     device: torch.device
-    cfg_parallel_ready: bool
+    guidance_parallel_ready: bool
     attention_kwargs: dict[str, Any] | None
     latent_num_frames: int
     latent_height: int
@@ -358,9 +358,9 @@ def step_denoised_latents(
         do_true_cfg=pipeline.do_classifier_free_guidance,
         per_request_scheduler=forward_ctx.video_audio_step_adapter,
     )
-    return pipeline._synchronize_cfg_parallel_step_output(
+    return pipeline._synchronize_guidance_parallel_step_output(
         latents,
-        do_true_cfg=pipeline.do_classifier_free_guidance,
+        guidance_parallel_ready=forward_ctx.guidance_parallel_ready,
     )
 
 
@@ -382,7 +382,7 @@ class LTXPhaseExecutor:
         prompt_context: LTXPromptContext | None = None,
     ) -> LTXPhaseResult:
         pipeline._check_forward_inputs(request_inputs, image=image)
-        cfg_parallel_ready = pipeline._setup_forward_runtime(req, request_inputs, attention_kwargs)
+        guidance_parallel_ready = pipeline._setup_forward_runtime(req, request_inputs, attention_kwargs)
         device = pipeline.device
         if prompt_context is None:
             prompt_context = pipeline._prepare_prompt_context(
@@ -429,7 +429,7 @@ class LTXPhaseExecutor:
             request_inputs=request_inputs,
             prompt_context=prompt_context,
             device=device,
-            cfg_parallel_ready=cfg_parallel_ready,
+            guidance_parallel_ready=guidance_parallel_ready,
             attention_kwargs=attention_kwargs,
             latent_num_frames=latent_num_frames,
             latent_height=latent_height,
@@ -449,7 +449,7 @@ class LTXPhaseExecutor:
             audio_coords=audio_coords,
             conditioning_mask=conditioning_mask,
         )
-        denoise_ctx = pipeline._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx)
+        denoise_ctx = pipeline._prepare_denoise_context_for_guidance(forward_ctx, denoise_ctx)
         state = LTXDenoiseExecutor.run(
             pipeline,
             LTXAVState(video=denoise_ctx.latents, audio=denoise_ctx.audio_latents),

vllm_omni/diffusion/models/ltx2/ltx2_guidance.py

@@ -10,12 +10,19 @@ from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
 import torch
+from vllm.logger import init_logger
 
 from vllm_omni.diffusion.distributed.parallel_state import (
-    get_cfg_group,
-    get_classifier_free_guidance_rank,
-    get_classifier_free_guidance_world_size,
+    get_cfg_group as get_guidance_parallel_group,
 )
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_classifier_free_guidance_rank as get_guidance_parallel_rank,
+)
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_classifier_free_guidance_world_size as get_guidance_parallel_world_size,
+)
+
+logger = init_logger(__name__)
 
 if TYPE_CHECKING:
     from .ltx2_denoise import LTXDenoiseContext, LTXForwardContext
@@ -110,10 +117,6 @@ class LTXGuidancePlan:
     def names(self) -> tuple[str, ...]:
         return tuple(item.name for item in self.passes)
 
-    @property
-    def cfg_parallel_compatible(self) -> bool:
-        return self.names == ("cond", "uncond") and not self.spec.do_rescale
-
 
 def x0_from_velocity(sample: torch.Tensor, velocity: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
     """Convert velocity to x0 using the official fp32 arithmetic order."""
@@ -212,23 +215,69 @@ class LTXGuidanceExecutor:
     """Build and execute official LTX multi-guidance Transformer passes."""
 
     @staticmethod
-    def validate_cfg_world_size(plan: LTXGuidancePlan, cfg_world_size: int) -> None:
-        if cfg_world_size == 1:
+    def validate_guidance_world_size(plan: LTXGuidancePlan, guidance_world_size: int) -> None:
+        del plan
+        if guidance_world_size < 1:
+            raise ValueError(f"LTX guidance parallelism requires a positive parallel size, got {guidance_world_size}.")
+
+    @staticmethod
+    def _parallel_assignments(pass_count: int, world_size: int) -> list[list[int]]:
+        assignments: list[list[int]] = [[] for _ in range(world_size)]
+        for pass_index in range(pass_count):
+            assignments[pass_index % world_size].append(pass_index)
+        return assignments
+
+    @classmethod
+    def _model_pass_count(
+        cls,
+        plan: LTXGuidancePlan,
+        guidance_parallel_ready: bool,
+        guidance_world_size: int,
+    ) -> int:
+        if not guidance_parallel_ready:
+            return len(plan.passes)
+        assignments = cls._parallel_assignments(len(plan.passes), guidance_world_size)
+        return max(len(indices) for indices in assignments)
+
+    @classmethod
+    def warn_if_imbalanced(
+        cls,
+        plan: LTXGuidancePlan,
+        guidance_world_size: int,
+        phase_name: str,
+    ) -> None:
+        pass_count = len(plan.passes)
+        if guidance_world_size <= 1 or pass_count % guidance_world_size == 0 or get_guidance_parallel_rank() != 0:
             return
-        if cfg_world_size != 2:
-            raise ValueError(f"LTX CFG parallelism supports cfg_parallel_size 1 or 2, got {cfg_world_size}.")
-        if not plan.cfg_parallel_compatible:
-            raise ValueError("LTX CFG parallelism only supports CFG-only guidance without rescale.")
+        slots_per_rank = math.ceil(pass_count / guidance_world_size)
+        total_slots = slots_per_rank * guidance_world_size
+        wasted_slots = total_slots - pass_count
+        utilization = 100.0 * pass_count / total_slots
+        logger.warning_once(
+            "LTX guidance parallelism is imbalanced for phase %r: %d guidance passes across %d ranks. "
+            "%d padded/discarded rank slot(s); expected guidance-slot utilization is %.1f%%.",
+            phase_name,
+            pass_count,
+            guidance_world_size,
+            wasted_slots,
+            utilization,
+        )
 
     @staticmethod
     def prepare_denoise_context(
         plan: LTXGuidancePlan,
-        cfg_parallel_ready: bool,
+        guidance_parallel_ready: bool,
+        guidance_world_size: int,
         denoise_ctx: LTXDenoiseContext,
     ) -> LTXDenoiseContext:
-        if len(plan.passes) > 1 and not cfg_parallel_ready:
-            denoise_ctx.video_coords = _repeat_batch(denoise_ctx.video_coords, len(plan.passes))
-            denoise_ctx.audio_coords = _repeat_batch(denoise_ctx.audio_coords, len(plan.passes))
+        model_pass_count = LTXGuidanceExecutor._model_pass_count(
+            plan,
+            guidance_parallel_ready,
+            guidance_world_size,
+        )
+        if model_pass_count > 1:
+            denoise_ctx.video_coords = _repeat_batch(denoise_ctx.video_coords, model_pass_count)
+            denoise_ctx.audio_coords = _repeat_batch(denoise_ctx.audio_coords, model_pass_count)
         return denoise_ctx
 
     @staticmethod
@@ -265,7 +314,30 @@ class LTXGuidanceExecutor:
         )
         return velocity_from_x0(sample, guided_x0, sigma)
 
-    def predict_parallel_cfg(
+    @staticmethod
+    def _guide_modality(
+        sample: torch.Tensor,
+        splits: dict[str, torch.Tensor],
+        sigma: torch.Tensor,
+        guidance: LTXModalityGuidance,
+        *,
+        model_sigma: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        model_sigma = sigma if model_sigma is None else model_sigma
+        # Official guidance reduces contiguous BSC tensors. The fused
+        # transformer output may be channel-major after splitting, which
+        # changes fp32 reduction order and can cross bf16 rounding bounds.
+        x0 = {name: x0_from_velocity(sample, value, model_sigma).contiguous() for name, value in splits.items()}
+        guided = combine_guided_x0(
+            cond=x0["cond"],
+            uncond_text=x0.get("uncond", 0.0),
+            uncond_perturbed=x0.get("ptb", 0.0),
+            uncond_modality=x0.get("mod", 0.0),
+            guidance=guidance,
+        )
+        return velocity_from_x0(sample, guided, sigma)
+
+    def predict_parallel_guidance(
         self,
         pipeline: Any,
         plan: LTXGuidancePlan,
@@ -275,54 +347,93 @@ class LTXGuidanceExecutor:
         forward_ctx: LTXForwardContext,
         denoise_ctx: LTXDenoiseContext,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        self.validate_cfg_world_size(plan, get_classifier_free_guidance_world_size())
+        guidance_world_size = get_guidance_parallel_world_size()
+        self.validate_guidance_world_size(plan, guidance_world_size)
+        guidance_rank = get_guidance_parallel_rank()
+        assignments = self._parallel_assignments(len(plan.passes), guidance_world_size)
+        local_pass_indices = assignments[guidance_rank]
+        model_pass_count = max(len(indices) for indices in assignments)
+
+        # Every rank executes and gathers the same number of slots. Shorter
+        # assignments use a conditional pass whose output is discarded.
+        padded_pass_indices: list[int | None] = local_pass_indices + [None] * (
+            model_pass_count - len(local_pass_indices)
+        )
+        local_passes = tuple(
+            plan.passes[0] if pass_index is None else plan.passes[pass_index] for pass_index in padded_pass_indices
+        )
+        local_plan = LTXGuidancePlan(spec=plan.spec, passes=local_passes)
+
         prompt = forward_ctx.prompt_context
-        negative = get_classifier_free_guidance_rank() == 1
-        video_context = (
-            prompt.negative_connector_prompt_embeds
-            if negative and plan.spec.video.do_cfg
-            else prompt.positive_connector_prompt_embeds
+        video_contexts: list[torch.Tensor] = []
+        audio_contexts: list[torch.Tensor] = []
+        for denoise_pass in local_passes:
+            video_context = (
+                prompt.negative_connector_prompt_embeds
+                if denoise_pass.negative_video_context
+                else prompt.positive_connector_prompt_embeds
+            )
+            audio_context = (
+                prompt.negative_connector_audio_prompt_embeds
+                if denoise_pass.negative_audio_context
+                else prompt.positive_connector_audio_prompt_embeds
+            )
+            if video_context is None or audio_context is None:
+                raise ValueError("Negative prompt context is required when LTX CFG is enabled.")
+            video_contexts.append(video_context)
+            audio_contexts.append(audio_context)
+
+        video_input = _repeat_batch(state.video, model_pass_count).to(prompt.positive_connector_prompt_embeds.dtype)
+        audio_input = _repeat_batch(state.audio, model_pass_count).to(
+            prompt.positive_connector_audio_prompt_embeds.dtype
         )
-        audio_context = (
-            prompt.negative_connector_audio_prompt_embeds
-            if negative and plan.spec.audio.do_cfg
-            else prompt.positive_connector_audio_prompt_embeds
-        )
-        if video_context is None or audio_context is None:
-            raise ValueError("Negative prompt context is required when LTX CFG is enabled.")
-        video_input = state.video.to(prompt.positive_connector_prompt_embeds.dtype)
-        audio_input = state.audio.to(prompt.positive_connector_audio_prompt_embeds.dtype)
+        attention_kwargs = dict(forward_ctx.attention_kwargs or {})
+        perturbations = build_perturbation_kwargs(local_plan, state.video.shape[0], video_input)
+        if perturbations:
+            attention_kwargs["ltx_perturbation_kwargs"] = perturbations
         kwargs = pipeline._build_transformer_kwargs(
             forward_ctx,
             denoise_ctx,
             hidden_states=video_input,
             audio_hidden_states=audio_input,
-            encoder_hidden_states=video_context,
-            audio_encoder_hidden_states=audio_context,
+            encoder_hidden_states=torch.cat(video_contexts),
+            audio_encoder_hidden_states=torch.cat(audio_contexts),
             encoder_attention_mask=None,
             audio_encoder_attention_mask=None,
             ts=timestep.expand(video_input.shape[0]),
+            attention_kwargs=attention_kwargs,
         )
-        local_video, local_audio = pipeline.predict_noise(**kwargs)
-        group = get_cfg_group()
-        video = group.all_gather(local_video, separate_tensors=True)
-        audio = group.all_gather(local_audio, separate_tensors=True)
+        with pipeline._transformer_cache_context("guided"):
+            local_video, local_audio = pipeline.transformer(**kwargs)
+        local_video_slots = local_video.chunk(model_pass_count)
+        local_audio_slots = local_audio.chunk(model_pass_count)
+
+        group = get_guidance_parallel_group()
+        gathered_video_slots = [group.all_gather(value, separate_tensors=True) for value in local_video_slots]
+        gathered_audio_slots = [group.all_gather(value, separate_tensors=True) for value in local_audio_slots]
+
+        video_splits: dict[str, torch.Tensor] = {}
+        audio_splits: dict[str, torch.Tensor] = {}
+        for pass_index, denoise_pass in enumerate(plan.passes):
+            owner_rank = pass_index % guidance_world_size
+            slot_index = pass_index // guidance_world_size
+            video_splits[denoise_pass.name] = gathered_video_slots[slot_index][owner_rank]
+            audio_splits[denoise_pass.name] = gathered_audio_slots[slot_index][owner_rank]
+
         video_sigma = pipeline.scheduler.sigmas[index]
         return (
-            self.combine_cfg_velocity(
+            self._guide_modality(
                 state.video,
-                video[0],
-                video[1],
+                video_splits,
                 video_sigma,
-                plan.spec.video.cfg_scale,
+                plan.spec.video,
                 model_sigma=pipeline._video_guidance_model_sigma(video_sigma, denoise_ctx),
             ),
-            self.combine_cfg_velocity(
+            self._guide_modality(
                 state.audio,
-                audio[0],
-                audio[1],
+                audio_splits,
                 forward_ctx.audio_scheduler.sigmas[index],
-                plan.spec.audio.cfg_scale,
+                plan.spec.audio,
             ),
         )
 
@@ -336,8 +447,8 @@ class LTXGuidanceExecutor:
         forward_ctx: LTXForwardContext,
         denoise_ctx: LTXDenoiseContext,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        if forward_ctx.cfg_parallel_ready:
-            return self.predict_parallel_cfg(pipeline, plan, index, timestep, state, forward_ctx, denoise_ctx)
+        if forward_ctx.guidance_parallel_ready:
+            return self.predict_parallel_guidance(pipeline, plan, index, timestep, state, forward_ctx, denoise_ctx)
 
         prompt = forward_ctx.prompt_context
         video_contexts: list[torch.Tensor] = []
@@ -382,30 +493,8 @@ class LTXGuidanceExecutor:
         video_splits = dict(zip(plan.names, video_velocity.chunk(pass_count), strict=True))
         audio_splits = dict(zip(plan.names, audio_velocity.chunk(pass_count), strict=True))
 
-        def guide_modality(
-            sample: torch.Tensor,
-            splits: dict[str, torch.Tensor],
-            sigma: torch.Tensor,
-            guidance: LTXModalityGuidance,
-            *,
-            model_sigma: torch.Tensor | None = None,
-        ) -> torch.Tensor:
-            model_sigma = sigma if model_sigma is None else model_sigma
-            # Official guidance reduces contiguous BSC tensors. The fused
-            # transformer output may be channel-major after splitting, which
-            # changes fp32 reduction order and can cross bf16 rounding bounds.
-            x0 = {name: x0_from_velocity(sample, value, model_sigma).contiguous() for name, value in splits.items()}
-            guided = combine_guided_x0(
-                cond=x0["cond"],
-                uncond_text=x0.get("uncond", 0.0),
-                uncond_perturbed=x0.get("ptb", 0.0),
-                uncond_modality=x0.get("mod", 0.0),
-                guidance=guidance,
-            )
-            return velocity_from_x0(sample, guided, sigma)
-
         return (
-            guide_modality(
+            self._guide_modality(
                 state.video,
                 video_splits,
                 pipeline.scheduler.sigmas[index],
@@ -415,7 +504,7 @@ class LTXGuidanceExecutor:
                     denoise_ctx,
                 ),
             ),
-            guide_modality(
+            self._guide_modality(
                 state.audio,
                 audio_splits,
                 forward_ctx.audio_scheduler.sigmas[index],

vllm_omni/diffusion/models/ltx2/ltx2_runtime.py

@@ -17,7 +17,9 @@ 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.parallel_state import get_classifier_free_guidance_world_size
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_classifier_free_guidance_world_size as get_guidance_parallel_world_size,
+)
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
@@ -44,8 +46,6 @@ from .ltx2_guidance import (
     LTX_GUIDANCE_EXECUTOR,
     LTXGuidanceExecutor,
     LTXGuidancePlan,
-    LTXGuidanceSpec,
-    LTXModalityGuidance,
 )
 from .ltx2_recipes import (
     LTXPhaseRecipe,
@@ -259,7 +259,7 @@ class LTXRuntime(
 
     def _enter_phase(self, phase: LTXPhaseRecipe) -> None:
         """Hook for a future phase-weight strategy."""
-        del phase
+        self._active_phase_name = phase.name
 
     def prepare_latents(
         self,
@@ -323,6 +323,10 @@ class LTXRuntime(
     def do_classifier_free_guidance(self):
         return self._guidance_plan.spec.do_cfg
 
+    @property
+    def do_guidance(self):
+        return len(self._guidance_plan.passes) > 1
+
     @property
     def interrupt(self):
         return self._interrupt
@@ -372,16 +376,16 @@ class LTXRuntime(
             ),
         )
 
-    def _synchronize_cfg_parallel_step_output(
+    def _synchronize_guidance_parallel_step_output(
         self,
         latents: tuple[torch.Tensor, torch.Tensor],
-        do_true_cfg: bool,
+        guidance_parallel_ready: bool,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        if not (do_true_cfg and get_classifier_free_guidance_world_size() > 1):
+        if not (guidance_parallel_ready and get_guidance_parallel_world_size() > 1):
             return latents
 
         # CUDA async execution otherwise permits numerical drift to accumulate
-        # across CFG-parallel denoise steps.
+        # across guidance-parallel denoise steps.
         latents = tuple(tensor.contiguous() for tensor in latents)
         current_omni_platform.synchronize()
         return latents
@@ -393,26 +397,16 @@ class LTXRuntime(
         attention_kwargs: dict[str, Any] | None,
     ) -> bool:
         self._guidance_plan = LTXGuidancePlan.build(request_inputs.guidance)
-        del attention_kwargs
+        del req, attention_kwargs
         self._interrupt = False
-        cfg_world_size = get_classifier_free_guidance_world_size()
-        if (
-            req.is_dummy_run()
-            and cfg_world_size == 2
-            and self.do_classifier_free_guidance
-            and not self._guidance_plan.cfg_parallel_compatible
-        ):
-            # The single-device warmup follows the full recipe. CFG2 executes
-            # the same per-modality CFG scales without STG/modality passes.
-            self._guidance_plan = LTXGuidancePlan.build(
-                LTXGuidanceSpec(
-                    video=LTXModalityGuidance(cfg_scale=self._guidance_plan.spec.video.cfg_scale),
-                    audio=LTXModalityGuidance(cfg_scale=self._guidance_plan.spec.audio.cfg_scale),
-                )
-            )
-        if self.do_classifier_free_guidance:
-            self.guidance_executor.validate_cfg_world_size(self._guidance_plan, cfg_world_size)
-        return self.do_classifier_free_guidance and cfg_world_size > 1
+        guidance_world_size = get_guidance_parallel_world_size()
+        self.guidance_executor.validate_guidance_world_size(self._guidance_plan, guidance_world_size)
+        self.guidance_executor.warn_if_imbalanced(
+            self._guidance_plan,
+            guidance_world_size,
+            getattr(self, "_active_phase_name", "generate"),
+        )
+        return self.do_guidance and guidance_world_size > 1
 
     def _check_forward_inputs(
         self,
@@ -594,14 +588,15 @@ class LTXRuntime(
         padded_length = latent_ops.get_sp_padded_audio_latent_length(requested_length, int(sp_size))
         return requested_length if provided_length in {requested_length, padded_length} else provided_length
 
-    def _prepare_denoise_context_for_cfg(
+    def _prepare_denoise_context_for_guidance(
         self,
         forward_ctx: LTXForwardContext,
         denoise_ctx: LTXDenoiseContext,
     ) -> LTXDenoiseContext:
         return self.guidance_executor.prepare_denoise_context(
             self._guidance_plan,
-            forward_ctx.cfg_parallel_ready,
+            forward_ctx.guidance_parallel_ready,
+            get_guidance_parallel_world_size(),
             denoise_ctx,
         )
 

vllm_omni/entrypoints/cli/serve.py

@@ -678,8 +678,7 @@ class OmniServeCommand(CLISubcommand):
             "--cfg-parallel-size",
             type=int,
             default=1,
-            choices=[1, 2],
-            help="Number of devices for CFG parallel computation for diffusion models. "
+            help="Number of devices used to execute diffusion guidance passes in parallel. "
             "Equivalent to setting DiffusionParallelConfig.cfg_parallel_size.",
         )
         omni_config_group.add_argument(