[diffusion] refactor: refactor cfg parallelism framework to support multi-branch CFG for LTX2 (#23736)

Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py

@@ -251,6 +251,24 @@ MODELS = {
         ],
     },
     # 13. Skill-only extra preset
+    "ltx23-two-stage-cfg-parallel": {
+        "path": "Lightricks/LTX-2.3",
+        "prompt": "A beautiful sunset over the ocean",
+        "negative_prompt": "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static.",
+        "seed": 1234,
+        "extra_args": [
+            "--pipeline-class-name=LTX2TwoStagePipeline",
+            "--width=1536",
+            "--height=1024",
+            "--num-frames=121",
+            "--fps=24",
+            "--num-inference-steps=30",
+            "--guidance-scale=3.0",
+            "--num-gpus=2",
+            "--cfg-parallel-size=2",
+        ],
+    },
+    # 14. Skill-only extra preset
     "hunyuanvideo": {
         "path": "hunyuanvideo-community/HunyuanVideo",
         "prompt": "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window.",
@@ -263,7 +281,7 @@ MODELS = {
             "--num-inference-steps=30",
         ],
     },
-    # 14. Skill-only extra preset
+    # 15. Skill-only extra preset
     # Requires: <repo>/inputs/diffusion_benchmark/figs/mova_single_person.jpg
     "mova-720p": {
         "path": "OpenMOSS-Team/MOVA-720p",
@@ -279,7 +297,7 @@ MODELS = {
             "--num-inference-steps=2",
         ],
     },
-    # 15. Skill-only extra preset
+    # 16. Skill-only extra preset
     "helios": {
         "path": "BestWishYsh/Helios-Base",
         "prompt": "A curious raccoon",

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -24,6 +24,7 @@ from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput
 from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config
 from sglang.multimodal_gen.configs.sample.sampling_params import DataType
 from sglang.multimodal_gen.configs.utils import update_config_from_args
+from sglang.multimodal_gen.runtime.distributed.cfg_policy import CFGPolicy
 from sglang.multimodal_gen.runtime.distributed.communication_op import (
     sequence_model_parallel_all_gather,
 )
@@ -207,6 +208,7 @@ class PipelineConfig:
     # controls the timestep embedding generation
     should_use_guidance: bool = True
     embedded_cfg_scale: float = 6.0
+    cfg_policy: CFGPolicy = field(default_factory=CFGPolicy)
     generator_device: str | None = None
     flow_shift: float | None = None
     disable_autocast: bool = False

python/sglang/multimodal_gen/runtime/distributed/cfg_parallel_utils.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+import dataclasses
+from typing import TYPE_CHECKING, Callable
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.distributed.cfg_policy import (
+    _apply_cfg_postprocess,
+    _unwrap,
+    _wrap,
+)
+from sglang.multimodal_gen.runtime.distributed.communication_op import (
+    cfg_model_parallel_all_gather,
+    cfg_model_parallel_all_reduce,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_cfg_group,
+    get_classifier_free_guidance_rank,
+    get_classifier_free_guidance_world_size,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.distributed.cfg_policy import (
+        CFGBranch,
+        CFGPolicy,
+    )
+
+# Tracks (n_branches, cfg_world_size, cfg_rank) tuples already logged so the
+# dispatch table is printed once per unique configuration, not once per step.
+_logged_dispatch_keys: set[tuple[int, int, int]] = set()
+
+
+def _run(
+    predict_fn: Callable[["CFGBranch"], "torch.Tensor | tuple[torch.Tensor, ...]"],
+    bid: int,
+    branches,
+) -> tuple[torch.Tensor, ...]:
+    branch = branches[bid]
+    device = get_local_torch_device()
+    local_branch = dataclasses.replace(
+        branch,
+        kwargs={
+            k: v.to(device) if isinstance(v, torch.Tensor) else v
+            for k, v in branch.kwargs.items()
+        },
+    )
+    raw = predict_fn(local_branch)
+    return _wrap(raw)
+
+
+def run_cfg_parallel(
+    policy: "CFGPolicy",
+    predict_fn: Callable[["CFGBranch"], "torch.Tensor | tuple[torch.Tensor, ...]"],
+) -> "list[torch.Tensor | tuple[torch.Tensor, ...]]":
+    """Dispatch CFG branches across ranks, all-gather results, return in branch order.
+
+    ``predict_fn`` is a closure capturing all step-varying state
+    (latent_model_input, timestep, model, etc.).  It is called with each
+    assigned ``CFGBranch`` and must return the raw ``_predict_noise`` output.
+
+    Idle ranks (cfg_world_size > n_branches) run branch 0 as a dummy forward
+    to obtain tensor shapes for the all-gather.
+
+    Returns a list indexed to match ``policy.branches``, identical on every rank.
+    """
+
+    cfg_rank = get_classifier_free_guidance_rank()
+    cfg_world_size = get_classifier_free_guidance_world_size()
+    branches = policy.branches
+    n_branches = len(branches)
+    assignments = dispatch_branches(n_branches, cfg_world_size)
+    branches_assigned_to_local_rank = assignments[cfg_rank]
+    max_num_branches_per_rank = max(len(a) for a in assignments)
+
+    if cfg_world_size > n_branches:
+        logger.warning_once(
+            "cfg_parallel_size=%d > n_branches=%d; %d GPU(s) will be idle for CFG",
+            cfg_world_size,
+            n_branches,
+            cfg_world_size - n_branches,
+        )
+
+    dispatch_key = (n_branches, cfg_world_size, cfg_rank)
+    if dispatch_key not in _logged_dispatch_keys:
+        _logged_dispatch_keys.add(dispatch_key)
+        branch_names = (
+            [branches[i].name for i in branches_assigned_to_local_rank]
+            if branches_assigned_to_local_rank
+            else ["(idle)"]
+        )
+        logger.info(
+            "CFG parallel dispatch: rank %d/%d -> [%s]",
+            cfg_rank,
+            cfg_world_size,
+            ", ".join(branch_names),
+        )
+
+    # perform the forward for local branches
+    predicts_from_local_branches: list[tuple[torch.Tensor, ...]] = [
+        _run(predict_fn, bid, branches) for bid in branches_assigned_to_local_rank
+    ]
+
+    if not predicts_from_local_branches:  # idle rank: run branch 0 for tensor shapes
+        predicts_from_local_branches.append(_run(predict_fn, 0, branches))
+
+    # pad the predicts to the length of max_num_branches_per_rank, to prepare for the all-gather later
+    ref = predicts_from_local_branches[0]
+    while len(predicts_from_local_branches) < max_num_branches_per_rank:
+        # TODO: cache this zero
+        predicts_from_local_branches.append(tuple(torch.zeros_like(t) for t in ref))
+
+    # All-gather each slot and output element with separate_tensors=True.
+    # all_slots[slot][elem] = list[Tensor] indexed by CFG rank; no reshape.
+    all_slots: list[list[list[torch.Tensor]]] = [
+        [
+            cfg_model_parallel_all_gather(p, dim=0, separate_tensors=True)
+            for p in slot_pred
+        ]
+        for slot_pred in predicts_from_local_branches
+    ]
+
+    # reorder the results in branch order: branch bid -> owner rank, slot.
+    n_elems = len(ref)
+    final: list[torch.Tensor | tuple[torch.Tensor, ...]] = []
+    for bid in range(n_branches):
+        owner = bid % cfg_world_size
+        slot = bid // cfg_world_size
+        elems = tuple(all_slots[slot][ei][owner] for ei in range(n_elems))
+        final.append(_unwrap(elems))
+    return final
+
+
+def run_two_branch_cfg_parallel(
+    policy: "CFGPolicy",
+    predict_fn: Callable[["CFGBranch"], "torch.Tensor | tuple[torch.Tensor, ...]"],
+    cfg_scale: float,
+    batch,
+    pipeline_config,
+) -> "torch.Tensor | tuple[torch.Tensor, ...]":
+    """Run standard two-pass CFG with the old all-reduce combine.
+
+    This keeps the existing WAN baselines: it avoids gathering both branch
+    predictions, and it preserves the bf16 arithmetic order used before the
+    multi-branch CFG dispatcher was added.
+    """
+
+    cfg_rank = get_classifier_free_guidance_rank()
+    pred_t = _run(predict_fn, cfg_rank, policy.branches)
+
+    if cfg_rank == 0:
+        partial = tuple(cfg_scale * p for p in pred_t)
+        cond_t = pred_t
+    else:
+        partial = tuple((1 - cfg_scale) * p for p in pred_t)
+        cond_t = tuple(torch.empty_like(p) for p in pred_t)
+
+    results = [cfg_model_parallel_all_reduce(p) for p in partial]
+    cond_t = tuple(get_cfg_group().broadcast(p, src=0) for p in cond_t)
+    results[0] = _apply_cfg_postprocess(results[0], cond_t[0], batch, pipeline_config)
+    return _unwrap(tuple(results))
+
+
+def dispatch_branches(n_branches: int, n_ranks: int) -> list[list[int]]:
+    """Assign branches to ranks in Round-robin fashion
+
+    Returns a list of length ``n_ranks`` where element ``r`` contains the
+    branch indices assigned to rank ``r``.  Branch ``i`` goes to rank
+    ``i % n_ranks``.
+
+    Example: 4 passes, 2 GPUs:
+        rank 0 -> [0, 2],  rank 1 -> [1, 3]
+    """
+    assignments: list[list[int]] = [[] for _ in range(n_ranks)]
+    for i in range(n_branches):
+        assignments[i % n_ranks].append(i)
+    return assignments

python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+
+
+@dataclass
+class CFGBranch:
+    """Immutable specification of one CFG branch forward pass.
+
+    Built once before the denoising loop; read-only across all steps.
+    """
+
+    name: str
+    is_conditional: bool
+    kwargs: dict[str, Any]
+
+    def configure_batch(self, batch: "Req") -> None:
+        """Set batch state before this branch's forward pass.
+
+        Override for richer per-branch context (e.g. a branch index instead of
+        a single boolean) when a model needs more than two guidance modes.
+        """
+        batch.is_cfg_negative = not self.is_conditional
+
+
+@dataclass
+class CFGPolicy:
+    """Owns the CFG branches for one generation run and combines their predictions.
+
+    Built once before the denoising loop via ``build()``, then used read-only
+    across all steps.  Subclass and override ``build()`` / ``combine()`` for
+    custom CFG schemes (N-branch, multi-output, etc.).
+
+    The default implementation handles standard 2-branch CFG.  With a single
+    branch (CFG disabled) ``combine()`` returns the prediction unchanged.
+    """
+
+    branches: list[CFGBranch] = field(default_factory=list)
+
+    def build(
+        self,
+        batch: "Req",
+        image_kwargs: dict[str, Any],
+        pos_cond_kwargs: dict[str, Any],
+        neg_cond_kwargs: dict[str, Any],
+    ) -> "CFGPolicy":
+        """Return a new policy with branches populated.
+
+        Called once before the denoising loop.  The returned policy is
+        immutable for the lifetime of the run.  Override to declare N branches.
+        """
+        branches = [CFGBranch("conditional", True, {**image_kwargs, **pos_cond_kwargs})]
+        if batch.do_classifier_free_guidance:
+            branches.append(
+                CFGBranch("unconditional", False, {**image_kwargs, **neg_cond_kwargs})
+            )
+        return dataclasses.replace(self, branches=branches)
+
+    def combine(
+        self,
+        predictions: list[torch.Tensor | tuple[torch.Tensor, ...]],
+        batch: "Req",
+        cfg_scale: float,
+        pipeline_config: Any,
+        *,
+        cfg_parallel: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+        """Combine branch predictions into the final noise estimate.
+
+        Default: standard 2-branch CFG formula applied element-wise, followed
+        by normalization / rescale / model-specific postprocess.
+        Single-branch (CFG disabled): returns the prediction unchanged.
+        Override for N-branch or multi-output models.
+        """
+        if len(predictions) == 1:
+            return predictions[0]
+        pos_t = _wrap(predictions[0])
+        neg_t = _wrap(predictions[1])
+        if cfg_parallel:
+            # Match the old CFG-parallel calculation: multiply the positive
+            # prediction by cfg_scale and the negative prediction by
+            # (1 - cfg_scale) before adding them. The serial CFG formula is
+            # mathematically equivalent, but bf16 rounding changes WAN outputs.
+            results = [
+                cfg_scale * p + (1 - cfg_scale) * n for p, n in zip(pos_t, neg_t)
+            ]
+        else:
+            results = [n + cfg_scale * (p - n) for p, n in zip(pos_t, neg_t)]
+        results[0] = _apply_cfg_postprocess(
+            results[0], pos_t[0], batch, pipeline_config
+        )
+        return _unwrap(tuple(results))
+
+
+# Helpers used by CFGPolicy and run_cfg_parallel.
+
+
+def _wrap(
+    pred: torch.Tensor | tuple[torch.Tensor, ...],
+) -> tuple[torch.Tensor, ...]:
+    return pred if isinstance(pred, tuple) else (pred,)
+
+
+def _unwrap(
+    pred: tuple[torch.Tensor, ...],
+) -> torch.Tensor | tuple[torch.Tensor, ...]:
+    return pred[0] if len(pred) == 1 else pred
+
+
+def _apply_cfg_postprocess(
+    noise_pred: torch.Tensor,
+    noise_pred_cond: torch.Tensor,
+    batch: "Req",
+    pipeline_config: Any,
+) -> torch.Tensor:
+    if batch.cfg_normalization and float(batch.cfg_normalization) > 0:
+        noise_pred = _apply_cfg_normalization(
+            noise_pred, noise_pred_cond, float(batch.cfg_normalization)
+        )
+    if batch.guidance_rescale > 0.0:
+        noise_pred = _rescale_noise_cfg(
+            noise_pred, noise_pred_cond, guidance_rescale=batch.guidance_rescale
+        )
+    return pipeline_config.postprocess_cfg_noise(batch, noise_pred, noise_pred_cond)
+
+
+def _apply_cfg_normalization(
+    noise_pred: torch.Tensor,
+    noise_pred_cond: torch.Tensor,
+    cfg_normalization: float,
+) -> torch.Tensor:
+    cond_f = noise_pred_cond.float()
+    pred_f = noise_pred.float()
+    ori_norm = torch.linalg.vector_norm(cond_f)
+    new_norm = torch.linalg.vector_norm(pred_f)
+    max_norm = ori_norm * cfg_normalization
+    if new_norm > max_norm:
+        noise_pred = noise_pred * (max_norm / new_norm)
+    return noise_pred
+
+
+def _rescale_noise_cfg(
+    noise_cfg: torch.Tensor,
+    noise_pred_text: torch.Tensor,
+    guidance_rescale: float = 0.0,
+) -> torch.Tensor:
+    std_text = noise_pred_text.std(
+        dim=list(range(1, noise_pred_text.ndim)), keepdim=True
+    )
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    return guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg

python/sglang/multimodal_gen/runtime/distributed/parallel_state.py

@@ -467,7 +467,7 @@ def get_dp_rank() -> int:
 def maybe_init_distributed_environment_and_model_parallel(
     tp_size: int,
     sp_size: int,
-    enable_cfg_parallel: bool,
+    cfg_degree: int = 1,
     ulysses_degree: int = 1,
     ring_degree: int = 1,
     dp_size: int = 1,
@@ -508,7 +508,7 @@ def maybe_init_distributed_environment_and_model_parallel(
     )
     initialize_model_parallel(
         data_parallel_size=dp_size,
-        classifier_free_guidance_degree=2 if enable_cfg_parallel else 1,
+        classifier_free_guidance_degree=cfg_degree,
         tensor_parallel_degree=tp_size,
         ulysses_degree=ulysses_degree,
         ring_degree=ring_degree,

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -127,7 +127,7 @@ class GPUWorker:
         # initialize the distributed environment
         maybe_init_distributed_environment_and_model_parallel(
             tp_size=self.server_args.tp_size,
-            enable_cfg_parallel=self.server_args.enable_cfg_parallel,
+            cfg_degree=self.server_args.cfg_parallel_degree or 1,
             ulysses_degree=self.server_args.ulysses_degree,
             ring_degree=self.server_args.ring_degree,
             sp_size=self.server_args.sp_degree,

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -37,7 +37,6 @@ from sglang.multimodal_gen.runtime.cache.cache_dit_integration import (
 )
 from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
 from sglang.multimodal_gen.runtime.distributed import (
-    cfg_model_parallel_all_reduce,
     get_local_torch_device,
     get_sp_group,
     get_sp_world_size,
@@ -45,12 +44,20 @@ from sglang.multimodal_gen.runtime.distributed import (
     get_world_group,
     get_world_size,
 )
+from sglang.multimodal_gen.runtime.distributed.cfg_parallel_utils import (
+    run_cfg_parallel,
+    run_two_branch_cfg_parallel,
+)
+from sglang.multimodal_gen.runtime.distributed.cfg_policy import (
+    CFGPolicy,
+    _unwrap,
+    _wrap,
+)
 from sglang.multimodal_gen.runtime.distributed.communication_op import (
     sequence_model_parallel_all_gather,
 )
 from sglang.multimodal_gen.runtime.distributed.parallel_state import (
-    get_cfg_group,
+    get_classifier_free_guidance_world_size,
-    get_classifier_free_guidance_rank,
 )
 from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
 from sglang.multimodal_gen.runtime.layers.attention.STA_configuration import (
@@ -118,6 +125,7 @@ class DenoisingContext:
     seq_len: int | None
     guidance: torch.Tensor
     is_warmup: bool
+    cfg_policy: CFGPolicy | None = None
     trajectory_timesteps: list[torch.Tensor] = field(default_factory=list)
     trajectory_latents: list[torch.Tensor] = field(default_factory=list)
     extra: dict[str, Any] = field(default_factory=dict)
@@ -733,6 +741,10 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         else:
             neg_cond_kwargs = {}
 
+        cfg_policy = server_args.pipeline_config.cfg_policy.build(
+            batch, image_kwargs, pos_cond_kwargs, neg_cond_kwargs
+        )
+
         return DenoisingContext(
             scheduler=scheduler,
             extra_step_kwargs=extra_step_kwargs,
@@ -751,6 +763,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
             seq_len=seq_len,
             guidance=guidance,
             is_warmup=batch.is_warmup,
+            cfg_policy=cfg_policy,
         )
 
     def _before_denoising_loop(
@@ -912,9 +925,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
             attn_metadata=step.attn_metadata,
             target_dtype=ctx.target_dtype,
             current_guidance_scale=step.current_guidance_scale,
-            image_kwargs=ctx.image_kwargs,
+            cfg_policy=ctx.cfg_policy,
-            pos_cond_kwargs=ctx.pos_cond_kwargs,
-            neg_cond_kwargs=ctx.neg_cond_kwargs,
             server_args=server_args,
             guidance=ctx.guidance,
             latents=ctx.latents,
@@ -1344,190 +1355,72 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         disable = local_rank != 0
         return tqdm(iterable=iterable, total=total, disable=disable)
 
-    def _rescale_noise_cfg(
+    def _predict_noise_with_cfg(
-        self, noise_cfg, noise_pred_text, guidance_rescale=0.0
-    ) -> torch.Tensor:
-        """
-        Rescale noise prediction according to guidance_rescale.
-
-        Based on findings of "Common Diffusion Noise Schedules and Sample Steps are Flawed"
-        (https://arxiv.org/pdf/2305.08891.pdf), Section 3.4.
-
-        Args:
-            noise_cfg: The noise prediction with guidance.
-            noise_pred_text: The text-conditioned noise prediction.
-            guidance_rescale: The guidance rescale factor.
-
-        Returns:
-            The rescaled noise prediction.
-        """
-        std_text = noise_pred_text.std(
-            dim=list(range(1, noise_pred_text.ndim)), keepdim=True
-        )
-        std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
-        # Rescale the results from guidance (fixes overexposure)
-        noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
-        # Mix with the original results from guidance by factor guidance_rescale
-        noise_cfg = (
-            guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
-        )
-        return noise_cfg
-
-    def _apply_cfg_normalization(
-        self,
-        noise_pred: torch.Tensor,
-        noise_pred_cond: torch.Tensor,
-        cfg_normalization: float,
-    ) -> torch.Tensor:
-        factor = float(cfg_normalization)
-        cond_f = noise_pred_cond.float()
-        pred_f = noise_pred.float()
-        ori_norm = torch.linalg.vector_norm(cond_f)
-        new_norm = torch.linalg.vector_norm(pred_f)
-        max_norm = ori_norm * factor
-
-        if new_norm > max_norm:
-            noise_pred = noise_pred * (max_norm / new_norm)
-        return noise_pred
-
-    def _apply_cfg_normalization_parallel(
-        self,
-        noise_pred: torch.Tensor,
-        noise_pred_cond: torch.Tensor | None,
-        cfg_normalization: float,
-        cfg_rank: int,
-    ) -> torch.Tensor:
-        # In cfg-parallel mode, only rank 0 has the conditional branch locally,
-        # so the reference norm has to be broadcast to the other ranks
-        factor = float(cfg_normalization)
-        pred_f = noise_pred.float()
-        new_norm = torch.linalg.vector_norm(pred_f)
-        if cfg_rank == 0:
-            assert noise_pred_cond is not None
-            ori_norm = torch.linalg.vector_norm(noise_pred_cond.float())
-        else:
-            ori_norm = torch.empty_like(new_norm)
-        ori_norm = get_cfg_group().broadcast(ori_norm, src=0)
-        max_norm = ori_norm * factor
-
-        if new_norm > max_norm:
-            noise_pred = noise_pred * (max_norm / new_norm)
-        return noise_pred
-
-    def _apply_guidance_rescale_parallel(
-        self,
-        noise_pred: torch.Tensor,
-        noise_pred_cond: torch.Tensor | None,
-        guidance_rescale: float,
-        cfg_rank: int,
-    ) -> torch.Tensor:
-        # Guidance rescale is still defined against the conditional branch, so
-        # cfg-parallel needs to broadcast that statistic to every rank
-        std_cfg = noise_pred.std(dim=list(range(1, noise_pred.ndim)), keepdim=True)
-        if cfg_rank == 0:
-            assert noise_pred_cond is not None
-            std_text = noise_pred_cond.std(
-                dim=list(range(1, noise_pred_cond.ndim)), keepdim=True
-            )
-        else:
-            std_text = torch.empty_like(std_cfg)
-        std_text = get_cfg_group().broadcast(std_text, src=0)
-        noise_pred_rescaled = noise_pred * (std_text / std_cfg)
-        return (
-            guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_pred
-        )
-
-    def _apply_model_specific_cfg_postprocess(
         self,
+        current_model: nn.Module,
+        latent_model_input: torch.Tensor,
+        timestep,
         batch: Req,
-        noise_pred: torch.Tensor,
+        timestep_index: int,
-        noise_pred_cond: torch.Tensor | None,
+        attn_metadata,
-        cfg_rank: int,
+        target_dtype,
-    ) -> torch.Tensor:
+        current_guidance_scale,
-        # keep model-specific CFG behavior out of the main denoising loop
+        cfg_policy: CFGPolicy,
-        # for cfg-parallel, broadcast cond noise first so the hook sees the same
+        server_args: ServerArgs,
-        # inputs as the serial path.
+        guidance: torch.Tensor,
-        if cfg_rank == 0:
+        latents: torch.Tensor,
-            assert noise_pred_cond is not None
+    ) -> "torch.Tensor | tuple[torch.Tensor, ...]":
-            cond_noise = noise_pred_cond
+        """Run all CFG branch forward passes and combine into the final noise estimate."""
+        cfg_scale = server_args.pipeline_config.get_classifier_free_guidance_scale(
+            batch, current_guidance_scale
+        )
+
+        def predict_fn(branch):
+            branch.configure_batch(batch)
+            with set_forward_context(
+                current_timestep=timestep_index,
+                attn_metadata=attn_metadata,
+                forward_batch=batch,
+            ):
+                raw = self._predict_noise(
+                    current_model=current_model,
+                    latent_model_input=latent_model_input,
+                    timestep=timestep,
+                    target_dtype=target_dtype,
+                    guidance=guidance,
+                    **branch.kwargs,
+                )
+            pred_t = _wrap(raw)
+            if len(pred_t) == 1:
+                pred_t = (
+                    server_args.pipeline_config.slice_noise_pred(pred_t[0], latents),
+                )
+            return _unwrap(pred_t)
+
+        if server_args.enable_cfg_parallel:
+            if (
+                len(cfg_policy.branches) == 2
+                and get_classifier_free_guidance_world_size() == 2
+            ):
+                return run_two_branch_cfg_parallel(
+                    cfg_policy,
+                    predict_fn,
+                    cfg_scale,
+                    batch,
+                    server_args.pipeline_config,
+                )
+            # perform cfg branches in parallel, following the cfg policy
+            predictions = run_cfg_parallel(cfg_policy, predict_fn)
         else:
-            # TODO: cache this?
+            # perform cfg branches one-by-one locally
-            cond_noise = torch.empty_like(noise_pred)
+            predictions = [predict_fn(branch) for branch in cfg_policy.branches]
-        cond_noise = get_cfg_group().broadcast(cond_noise, src=0)
 
-        # qwen-image uses true_cfg_scale, match the per-token norm back to the conditional branch
+        return cfg_policy.combine(
-        return self.server_args.pipeline_config.postprocess_cfg_noise(
+            predictions,
-            batch, noise_pred, cond_noise
+            batch,
-        )
+            cfg_scale,
-
+            server_args.pipeline_config,
-    def _combine_cfg_parallel(
+            cfg_parallel=server_args.enable_cfg_parallel,
-        self,
-        batch: Req,
-        noise_pred_cond: torch.Tensor | None,
-        noise_pred_uncond: torch.Tensor | None,
-        cfg_scale: float,
-        cfg_rank: int,
-    ) -> torch.Tensor:
-        # cfg-parallel splits cond / uncond across ranks and reconstructs the
-        # final CFG result with an all-reduce.
-        if cfg_rank == 0:
-            assert noise_pred_cond is not None
-            partial = cfg_scale * noise_pred_cond
-        else:
-            assert noise_pred_uncond is not None
-            partial = (1 - cfg_scale) * noise_pred_uncond
-
-        noise_pred = cfg_model_parallel_all_reduce(partial)
-
-        if batch.cfg_normalization and float(batch.cfg_normalization) > 0:
-            noise_pred = self._apply_cfg_normalization_parallel(
-                noise_pred,
-                noise_pred_cond,
-                batch.cfg_normalization,
-                cfg_rank,
-            )
-
-        if batch.guidance_rescale > 0.0:
-            noise_pred = self._apply_guidance_rescale_parallel(
-                noise_pred,
-                noise_pred_cond,
-                batch.guidance_rescale,
-                cfg_rank,
-            )
-
-        return self._apply_model_specific_cfg_postprocess(
-            batch, noise_pred, noise_pred_cond, cfg_rank
-        )
-
-    def _combine_cfg_serial(
-        self,
-        batch: Req,
-        noise_pred_cond: torch.Tensor,
-        noise_pred_uncond: torch.Tensor,
-        cfg_scale: float,
-    ) -> torch.Tensor:
-        # Serial CFG keeps both branches local and is the reference path that
-        # model-specific postprocessing hooks should match.
-        noise_pred = noise_pred_uncond + cfg_scale * (
-            noise_pred_cond - noise_pred_uncond
-        )
-
-        if batch.cfg_normalization and float(batch.cfg_normalization) > 0:
-            noise_pred = self._apply_cfg_normalization(
-                noise_pred,
-                noise_pred_cond,
-                batch.cfg_normalization,
-            )
-
-        if batch.guidance_rescale > 0.0:
-            noise_pred = self._rescale_noise_cfg(
-                noise_pred,
-                noise_pred_cond,
-                guidance_rescale=batch.guidance_rescale,
-            )
-
-        return self.server_args.pipeline_config.postprocess_cfg_noise(
-            batch, noise_pred, noise_pred_cond
         )
 
     def _build_attn_metadata(
@@ -1691,110 +1584,6 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
             **kwargs,
         )
 
-    def _predict_noise_with_cfg(
-        self,
-        current_model: nn.Module,
-        latent_model_input: torch.Tensor,
-        timestep,
-        batch: Req,
-        timestep_index: int,
-        attn_metadata,
-        target_dtype,
-        current_guidance_scale,
-        image_kwargs: dict[str, Any],
-        pos_cond_kwargs: dict[str, Any],
-        neg_cond_kwargs: dict[str, Any],
-        server_args,
-        guidance,
-        latents,
-    ):
-        """
-        Predict the noise residual with classifier-free guidance.
-
-        Args:
-            current_model: The transformer model to use for the current step.
-            latent_model_input: The input latents for the model.
-            timestep: The expanded timestep tensor.
-            batch: The current batch information.
-            timestep_index: The current timestep index.
-            attn_metadata: Attention metadata for custom backends.
-            target_dtype: The target data type for autocasting.
-            current_guidance_scale: The guidance scale for the current step.
-            image_kwargs: Keyword arguments for image conditioning.
-            pos_cond_kwargs: Keyword arguments for positive prompt conditioning.
-            neg_cond_kwargs: Keyword arguments for negative prompt conditioning.
-
-        Returns:
-            The predicted noise.
-        """
-        noise_pred_cond: torch.Tensor | None = None
-        noise_pred_uncond: torch.Tensor | None = None
-        cfg_rank = get_classifier_free_guidance_rank()
-        # positive pass
-        if not (server_args.enable_cfg_parallel and cfg_rank != 0):
-            batch.is_cfg_negative = False
-            with set_forward_context(
-                current_timestep=timestep_index,
-                attn_metadata=attn_metadata,
-                forward_batch=batch,
-            ):
-                noise_pred_cond = self._predict_noise(
-                    current_model=current_model,
-                    latent_model_input=latent_model_input,
-                    timestep=timestep,
-                    target_dtype=target_dtype,
-                    guidance=guidance,
-                    **image_kwargs,
-                    **pos_cond_kwargs,
-                )
-                # TODO: can it be moved to after _predict_noise_with_cfg?
-                noise_pred_cond = server_args.pipeline_config.slice_noise_pred(
-                    noise_pred_cond, latents
-                )
-        if not batch.do_classifier_free_guidance:
-            return noise_pred_cond
-
-        # negative pass
-        if not server_args.enable_cfg_parallel or cfg_rank != 0:
-            batch.is_cfg_negative = True
-            with set_forward_context(
-                current_timestep=timestep_index,
-                attn_metadata=attn_metadata,
-                forward_batch=batch,
-            ):
-                noise_pred_uncond = self._predict_noise(
-                    current_model=current_model,
-                    latent_model_input=latent_model_input,
-                    timestep=timestep,
-                    target_dtype=target_dtype,
-                    guidance=guidance,
-                    **image_kwargs,
-                    **neg_cond_kwargs,
-                )
-                noise_pred_uncond = server_args.pipeline_config.slice_noise_pred(
-                    noise_pred_uncond, latents
-                )
-        cfg_scale = server_args.pipeline_config.get_classifier_free_guidance_scale(
-            batch, current_guidance_scale
-        )
-
-        if server_args.enable_cfg_parallel:
-            return self._combine_cfg_parallel(
-                batch,
-                noise_pred_cond,
-                noise_pred_uncond,
-                cfg_scale,
-                cfg_rank,
-            )
-
-        assert noise_pred_cond is not None and noise_pred_uncond is not None
-        return self._combine_cfg_serial(
-            batch,
-            noise_pred_cond,
-            noise_pred_uncond,
-            cfg_scale,
-        )
-
     def prepare_sta_param(self, batch: Req, server_args: ServerArgs):
         """
         Prepare Sliding Tile Attention (STA) parameters and settings.

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py

@@ -7,6 +7,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils impo
     clone_scheduler_runtime,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
+    StageParallelismType,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import (
     LTX2DenoisingStage,
 )
@@ -138,6 +141,14 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
             )
         ]
 
+    @property
+    def parallelism_type(self) -> StageParallelismType:
+        # Stage 2 is distilled and always runs with CFG disabled, so non-main
+        # CFG ranks should wait at a barrier rather than run a redundant forward.
+        if self.server_args.enable_cfg_parallel:
+            return StageParallelismType.MAIN_RANK_ONLY
+        return StageParallelismType.REPLICATED
+
     @staticmethod
     def _randn_like_with_batch_generators(
         reference_tensor: torch.Tensor, batch: Req

python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py

@@ -8,12 +8,29 @@ from diffusers.utils.torch_utils import randn_tensor
 from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
     is_ltx23_native_variant,
 )
-from sglang.multimodal_gen.runtime.distributed import get_sp_world_size
+from sglang.multimodal_gen.runtime.distributed import (
+    get_local_torch_device,
+    get_sp_world_size,
+)
+from sglang.multimodal_gen.runtime.distributed.cfg_parallel_utils import (
+    dispatch_branches,
+)
+from sglang.multimodal_gen.runtime.distributed.communication_op import (
+    cfg_model_parallel_all_gather,
+    cfg_model_parallel_all_reduce,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_classifier_free_guidance_rank,
+    get_classifier_free_guidance_world_size,
+)
 from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
 from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import (
     clone_scheduler_runtime,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
+    StageParallelismType,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import (
     DenoisingContext,
     DenoisingStage,
@@ -144,6 +161,179 @@ class LTX2DenoisingStage(DenoisingStage):
             dtype=reference_tensor.dtype,
         )
 
+    @property
+    def parallelism_type(self) -> StageParallelismType:
+        if self.server_args.enable_cfg_parallel:
+            return StageParallelismType.CFG_PARALLEL
+        return StageParallelismType.REPLICATED
+
+    @staticmethod
+    def _combine_cfg_parallel_av(
+        video: torch.Tensor,
+        audio: torch.Tensor,
+        guidance_scale: float,
+        cfg_rank: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """All-reduce video and audio predictions across CFG ranks.
+
+        Rank 0 (cond) contributes ``guidance_scale * pred``.
+        Rank 1 (uncond) contributes ``(1 - guidance_scale) * pred``.
+        Higher CFG ranks, if configured for multi-pass guidance, contribute
+        zeros on the two-branch path.
+        The sum reconstructs ``uncond + guidance_scale * (cond - uncond)``.
+        """
+        if cfg_rank == 0:
+            video_partial = guidance_scale * video
+            audio_partial = guidance_scale * audio
+        elif cfg_rank == 1:
+            video_partial = (1.0 - guidance_scale) * video
+            audio_partial = (1.0 - guidance_scale) * audio
+        else:
+            video_partial = torch.zeros_like(video)
+            audio_partial = torch.zeros_like(audio)
+        return (
+            cfg_model_parallel_all_reduce(video_partial),
+            cfg_model_parallel_all_reduce(audio_partial),
+        )
+
+    def _run_legacy_one_stage_multi_branch_cfg_parallel(
+        self,
+        *,
+        base_model_kwargs: dict[str, object],
+        ctx: "LTX2DenoisingContext",
+        step: "DenoisingStepState",
+        encoder_hidden_states: torch.Tensor,
+        audio_encoder_hidden_states: torch.Tensor,
+        encoder_attention_mask: torch.Tensor | None,
+        negative_encoder_hidden_states: torch.Tensor,
+        negative_audio_encoder_hidden_states: torch.Tensor,
+        negative_encoder_attention_mask: torch.Tensor | None,
+        need_perturbed: bool,
+        need_modality: bool,
+        stage1_guider_params: dict[str, object],
+    ) -> dict[str, tuple[torch.Tensor, torch.Tensor]]:
+        """Multi-branch CFG parallel for the legacy LTX-2.3 one-stage path.
+
+        Distributes up to 4 forward passes (cond, neg, perturbed, modality)
+        across CFG ranks via round-robin.  Each rank runs only its assigned
+        passes, then an all-gather collects every output so all ranks can
+        compute the guidance combination locally.
+        """
+        cfg_rank = get_classifier_free_guidance_rank()
+        cfg_world_size = get_classifier_free_guidance_world_size()
+
+        # Build kwargs for every pass in canonical order.
+        all_passes: list[tuple[str, dict[str, object]]] = [
+            (
+                "cond",
+                self._build_ltx2_model_kwargs(
+                    ctx,
+                    base_model_kwargs,
+                    encoder_hidden_states=encoder_hidden_states,
+                    audio_encoder_hidden_states=audio_encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                ),
+            ),
+            (
+                "neg",
+                self._build_ltx2_model_kwargs(
+                    ctx,
+                    base_model_kwargs,
+                    encoder_hidden_states=negative_encoder_hidden_states,
+                    audio_encoder_hidden_states=negative_audio_encoder_hidden_states,
+                    encoder_attention_mask=negative_encoder_attention_mask,
+                ),
+            ),
+        ]
+        if need_perturbed:
+            all_passes.append(
+                (
+                    "perturbed",
+                    self._build_ltx2_model_kwargs(
+                        ctx,
+                        base_model_kwargs,
+                        encoder_hidden_states=encoder_hidden_states,
+                        audio_encoder_hidden_states=audio_encoder_hidden_states,
+                        encoder_attention_mask=encoder_attention_mask,
+                        skip_video_self_attn_blocks=tuple(
+                            stage1_guider_params["video_stg_blocks"]
+                        ),
+                        skip_audio_self_attn_blocks=tuple(
+                            stage1_guider_params["audio_stg_blocks"]
+                        ),
+                    ),
+                )
+            )
+        if need_modality:
+            all_passes.append(
+                (
+                    "modality",
+                    self._build_ltx2_model_kwargs(
+                        ctx,
+                        base_model_kwargs,
+                        encoder_hidden_states=encoder_hidden_states,
+                        audio_encoder_hidden_states=audio_encoder_hidden_states,
+                        encoder_attention_mask=encoder_attention_mask,
+                        disable_a2v_cross_attn=True,
+                        disable_v2a_cross_attn=True,
+                    ),
+                )
+            )
+
+        pass_names = [name for name, _ in all_passes]
+        n_passes = len(pass_names)
+        assignments = dispatch_branches(n_passes, cfg_world_size)
+        my_indices = assignments[cfg_rank]
+        max_local = max(len(a) for a in assignments)
+
+        local_videos: list[torch.Tensor] = []
+        local_audios: list[torch.Tensor] = []
+
+        indices_to_run = my_indices if my_indices else [0]
+        with set_forward_context(
+            current_timestep=step.step_index, attn_metadata=step.attn_metadata
+        ):
+            for idx in indices_to_run:
+                _, kwargs = all_passes[idx]
+                v, a = step.current_model(**kwargs)
+                local_videos.append(v.float())
+                local_audios.append(a.float())
+
+        if not my_indices:
+            # This rank has no real branch, but it still needs tensor shapes for all-gather.
+            # The dummy branch above provides the shapes; zeros keep this rank from contributing.
+            local_videos = [torch.zeros_like(local_videos[0])]
+            local_audios = [torch.zeros_like(local_audios[0])]
+
+        # Pad to max_local for unbalanced cases (n_passes not divisible by n_ranks).
+        while len(local_videos) < max_local:
+            local_videos.append(torch.zeros_like(local_videos[0]))
+            local_audios.append(torch.zeros_like(local_audios[0]))
+
+        # Stack -> [max_local, B, ...], flatten to [max_local*B, ...] for all-gather.
+        local_v = torch.stack(local_videos, dim=0)
+        local_a = torch.stack(local_audios, dim=0)
+        B = local_v.shape[1]
+        local_v_flat = local_v.reshape(max_local * B, *local_v.shape[2:])
+        local_a_flat = local_a.reshape(max_local * B, *local_a.shape[2:])
+
+        # All-gather along batch dim -> [cfg_world_size * max_local * B, ...].
+        all_v_flat = cfg_model_parallel_all_gather(local_v_flat, dim=0)
+        all_a_flat = cfg_model_parallel_all_gather(local_a_flat, dim=0)
+
+        # Reshape to [cfg_world_size, max_local, B, ...].
+        all_v = all_v_flat.reshape(cfg_world_size, max_local, B, *all_v_flat.shape[1:])
+        all_a = all_a_flat.reshape(cfg_world_size, max_local, B, *all_a_flat.shape[1:])
+
+        # Branch i was run by rank (i % cfg_world_size) at slot (i // cfg_world_size).
+        return {
+            name: (
+                all_v[i % cfg_world_size, i // cfg_world_size],
+                all_a[i % cfg_world_size, i // cfg_world_size],
+            )
+            for i, name in enumerate(pass_names)
+        }
+
     @staticmethod
     def _get_video_latent_num_frames_for_model(
         batch: Req, server_args: ServerArgs, latents: torch.Tensor
@@ -200,7 +390,7 @@ class LTX2DenoisingStage(DenoisingStage):
         return latents[:, :orig_s, :].contiguous()
 
     def _maybe_enable_cache_dit(self, num_inference_steps: int, batch: Req) -> None:
-        """Disable cache-dit for TI2V-style requests (image-conditioned), to avoid stale activations.
+        """Disable cache-dit for TI2V-style requests to avoid stale activations.
 
         NOTE: base denoising stage calls this hook with (num_inference_steps, batch).
         """
@@ -231,6 +421,47 @@ class LTX2DenoisingStage(DenoisingStage):
         factor = rescale_scale * factor + (1.0 - rescale_scale)
         return pred * factor
 
+    @classmethod
+    def _ltx2_combine_guided_x0_parallel(
+        cls,
+        *,
+        latents: torch.Tensor,
+        local_velocities: dict[str, torch.Tensor],
+        sigma: float | torch.Tensor,
+        cfg_scale: float,
+        stg_scale: float,
+        rescale_scale: float,
+        modality_scale: float,
+    ) -> torch.Tensor:
+        """Combine stage-1 guidance passes that were split across CFG ranks.
+
+        Each pass is one model forward with a different conditioning setup:
+        positive prompt, negative prompt, attention-disabled perturbation, or
+        audio/video cross-attention disabled. A rank only owns some passes, so
+        it contributes weighted x0 terms for those passes and all-reduce
+        reconstructs the full guided x0 on every rank.
+        """
+        coefficients = {
+            "cond": cfg_scale + stg_scale + modality_scale - 1.0,
+            "neg": 1.0 - cfg_scale,
+            "perturbed": -stg_scale,
+            "modality": 1.0 - modality_scale,
+        }
+        first_velocity = next(iter(local_velocities.values()))
+        template = cls._ltx2_velocity_to_x0(latents, first_velocity, sigma)
+        cond_partial = torch.zeros_like(template)
+        pred_partial = torch.zeros_like(template)
+
+        for name, velocity in local_velocities.items():
+            denoised = cls._ltx2_velocity_to_x0(latents, velocity, sigma)
+            if name == "cond":
+                cond_partial = cond_partial + denoised
+            pred_partial = pred_partial + denoised * coefficients[name]
+
+        cond = cfg_model_parallel_all_reduce(cond_partial)
+        pred = cfg_model_parallel_all_reduce(pred_partial)
+        return cls._ltx2_apply_rescale(cond, pred, rescale_scale)
+
     @staticmethod
     def _ltx2_channelwise_normalize(noise: torch.Tensor) -> torch.Tensor:
         return noise.sub_(noise.mean(dim=(-2, -1), keepdim=True)).div_(
@@ -274,6 +505,8 @@ class LTX2DenoisingStage(DenoisingStage):
         ctx: LTX2DenoisingContext,
         *,
         substep: bool,
+        batch: Req | None = None,
+        is_audio: bool = False,
     ) -> torch.Tensor:
         generator = (
             ctx.res2s_substep_noise_generator
@@ -282,7 +515,51 @@ class LTX2DenoisingStage(DenoisingStage):
         )
         if generator is None:
             raise ValueError("LTX-2 res2s noise generator was not initialized.")
-        return cls._ltx2_res2s_new_noise(reference_tensor, generator)
+        if batch is not None and get_sp_world_size() > 1 and reference_tensor.ndim == 3:
+            full_shape = (
+                getattr(batch, "raw_audio_latent_shape", None)
+                if is_audio
+                else getattr(batch, "raw_latent_shape", None)
+            )
+            did_shard = (
+                getattr(batch, "did_sp_shard_audio_latents", False)
+                if is_audio
+                else getattr(batch, "did_sp_shard_latents", False)
+            )
+            if full_shape is not None and did_shard:
+                # HQ res2s normalizes SDE noise over the complete latent. If
+                # each SP rank normalizes only its local slice, the sampler
+                # follows a different trajectory. Recreate the same full noise
+                # on every rank, then keep the time slice owned by this rank.
+                full_noise = cls._ltx2_res2s_new_noise(
+                    torch.empty(
+                        tuple(int(dim) for dim in full_shape),
+                        device=reference_tensor.device,
+                        dtype=reference_tensor.dtype,
+                    ),
+                    generator,
+                )
+                if is_audio:
+                    start = int(batch.sp_audio_start_frame)
+                    end = start + int(batch.sp_audio_latent_num_frames)
+                else:
+                    start = int(batch.sp_video_start_frame) * int(
+                        batch.sp_video_tokens_per_frame
+                    )
+                    end = start + int(reference_tensor.shape[1])
+                sliced = full_noise[:, start : min(end, int(full_noise.shape[1])), :]
+                if int(sliced.shape[1]) < int(reference_tensor.shape[1]):
+                    pad_len = int(reference_tensor.shape[1]) - int(sliced.shape[1])
+                    pad = torch.zeros(
+                        (sliced.shape[0], pad_len, sliced.shape[2]),
+                        device=sliced.device,
+                        dtype=sliced.dtype,
+                    )
+                    sliced = torch.cat([sliced, pad], dim=1)
+                return sliced.to(dtype=reference_tensor.dtype)
+        return cls._ltx2_res2s_new_noise(reference_tensor, generator).to(
+            dtype=reference_tensor.dtype
+        )
 
     @staticmethod
     def _ltx2_apply_clean_latent_mask(
@@ -480,12 +757,16 @@ class LTX2DenoisingStage(DenoisingStage):
         midpoint_audio_det = anchor_audio + h * a21 * eps1_audio
 
         sub_noise_video = (
-            self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=True)
+            self._ltx2_res2s_noise_like(
+                ctx.latents, ctx, substep=True, batch=batch
+            ).float()
             if ctx.use_native_hq_res2s_sde_noise
             else self._randn_like_with_batch_generators(ctx.latents, batch).float()
         )
         sub_noise_audio = (
-            self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=True)
+            self._ltx2_res2s_noise_like(
+                ctx.audio_latents, ctx, substep=True, batch=batch, is_audio=True
+            ).float()
             if ctx.use_native_hq_res2s_sde_noise
             else self._randn_like_with_batch_generators(
                 ctx.audio_latents, batch
@@ -552,12 +833,16 @@ class LTX2DenoisingStage(DenoisingStage):
         next_audio_det = anchor_audio + h * (b1 * eps1_audio + b2 * eps2_audio)
 
         step_noise_video = (
-            self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=False)
+            self._ltx2_res2s_noise_like(
+                ctx.latents, ctx, substep=False, batch=batch
+            ).float()
             if ctx.use_native_hq_res2s_sde_noise
             else self._randn_like_with_batch_generators(ctx.latents, batch).float()
         )
         step_noise_audio = (
-            self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=False)
+            self._ltx2_res2s_noise_like(
+                ctx.audio_latents, ctx, substep=False, batch=batch, is_audio=True
+            ).float()
             if ctx.use_native_hq_res2s_sde_noise
             else self._randn_like_with_batch_generators(
                 ctx.audio_latents, batch
@@ -746,7 +1031,8 @@ class LTX2DenoisingStage(DenoisingStage):
             return tensor
         if tensor.shape[0] <= 0 or int(target_batch_size) % int(tensor.shape[0]) != 0:
             raise ValueError(
-                f"Cannot repeat tensor with batch={tensor.shape[0]} to target_batch_size={target_batch_size}"
+                "Cannot repeat tensor with batch="
+                f"{tensor.shape[0]} to target_batch_size={target_batch_size}"
             )
         repeat_factor = int(target_batch_size) // int(tensor.shape[0])
         return tensor.repeat(repeat_factor, *([1] * (tensor.ndim - 1)))
@@ -1048,6 +1334,12 @@ class LTX2DenoisingStage(DenoisingStage):
             kwargs["disable_v2a_cross_attn"] = True
         if perturbation_configs is not None:
             kwargs["perturbation_configs"] = perturbation_configs
+        if self.server_args.enable_cfg_parallel:
+            device = get_local_torch_device()
+            return {
+                k: v.to(device) if isinstance(v, torch.Tensor) else v
+                for k, v in kwargs.items()
+            }
         return kwargs
 
     @staticmethod
@@ -1269,6 +1561,23 @@ class LTX2DenoisingStage(DenoisingStage):
                     clean_latent_background=clean_latent_background,
                 )
             )
+
+        # Batch tensors are broadcast from CFG rank 0 and remain on its device.
+        # Move every context tensor that will be used in model forward passes or
+        # scheduler steps to the local device once here, before the loop begins.
+        if server_args.enable_cfg_parallel:
+            device = get_local_torch_device()
+            ctx.latents = ctx.latents.to(device)
+            ctx.timesteps = ctx.timesteps.to(device)
+            if ctx.audio_latents is not None:
+                ctx.audio_latents = ctx.audio_latents.to(device)
+            if ctx.guidance is not None:
+                ctx.guidance = ctx.guidance.to(device)
+            if ctx.denoise_mask is not None:
+                ctx.denoise_mask = ctx.denoise_mask.to(device)
+            if ctx.clean_latent is not None:
+                ctx.clean_latent = ctx.clean_latent.to(device)
+
         return ctx
 
     def _before_denoising_loop(
@@ -1346,6 +1655,38 @@ class LTX2DenoisingStage(DenoisingStage):
         )
         use_official_cfg_path = stage1_guider_params is None
         if use_official_cfg_path:
+            cfg_parallel = (
+                server_args.enable_cfg_parallel and batch.do_classifier_free_guidance
+            )
+            cfg_rank = get_classifier_free_guidance_rank() if cfg_parallel else 0
+
+            if cfg_parallel:
+                if cfg_rank == 0:
+                    encoder_hidden_states = batch.prompt_embeds[0]
+                    audio_encoder_hidden_states = batch.audio_prompt_embeds[0]
+                    encoder_attention_mask = prompt_attention_mask
+                elif cfg_rank == 1:
+                    encoder_hidden_states = batch.negative_prompt_embeds[0]
+                    audio_encoder_hidden_states = batch.negative_audio_prompt_embeds[0]
+                    encoder_attention_mask = self._get_ltx_prompt_attention_mask(
+                        batch,
+                        is_ltx23_variant=(
+                            ctx.is_ltx23_variant and not ctx.use_ltx23_legacy_one_stage
+                        ),
+                        negative=True,
+                    )
+                else:
+                    encoder_hidden_states = batch.prompt_embeds[0]
+                    audio_encoder_hidden_states = batch.audio_prompt_embeds[0]
+                    encoder_attention_mask = prompt_attention_mask
+                model_kwargs = self._build_ltx2_model_kwargs(
+                    ctx,
+                    base_model_kwargs,
+                    encoder_hidden_states=encoder_hidden_states,
+                    audio_encoder_hidden_states=audio_encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                )
+            else:
                 model_kwargs = self._build_ltx2_model_kwargs(
                     ctx,
                     base_model_kwargs,
@@ -1359,7 +1700,8 @@ class LTX2DenoisingStage(DenoisingStage):
                         model_kwargs, cfg_batch_size
                     )
                     model_kwargs["encoder_hidden_states"] = torch.cat(
-                    [batch.negative_prompt_embeds[0], batch.prompt_embeds[0]], dim=0
+                        [batch.negative_prompt_embeds[0], batch.prompt_embeds[0]],
+                        dim=0,
                     )
                     model_kwargs["audio_encoder_hidden_states"] = torch.cat(
                         [
@@ -1392,7 +1734,11 @@ class LTX2DenoisingStage(DenoisingStage):
 
             model_video = model_video.float()
             model_audio = model_audio.float()
-            if batch.do_classifier_free_guidance:
+            if cfg_parallel:
+                model_video, model_audio = self._combine_cfg_parallel_av(
+                    model_video, model_audio, float(batch.guidance_scale), cfg_rank
+                )
+            elif batch.do_classifier_free_guidance:
                 model_video_uncond, model_video_text = model_video.chunk(2)
                 model_audio_uncond, model_audio_text = model_audio.chunk(2)
                 model_video = model_video_uncond + (
@@ -1554,6 +1900,15 @@ class LTX2DenoisingStage(DenoisingStage):
             float(stage1_guider_params["video_modality_scale"]) != 1.0
             or float(stage1_guider_params["audio_modality_scale"]) != 1.0
         )
+        stage1_cfg_parallel = (
+            server_args.enable_cfg_parallel and not ctx.use_ltx23_legacy_one_stage
+        )
+        stage1_cfg_rank = (
+            get_classifier_free_guidance_rank() if stage1_cfg_parallel else 0
+        )
+        stage1_cfg_world_size = (
+            get_classifier_free_guidance_world_size() if stage1_cfg_parallel else 1
+        )
         # NOTE: this flag must be identical across all SP ranks so that every
         # rank executes the same number of model-forward calls (each of which
         # contains NCCL collectives).
@@ -1722,8 +2077,21 @@ class LTX2DenoisingStage(DenoisingStage):
                             )
                         )
 
-                    num_passes = len(pass_specs)
+                    execution_pass_specs = (
-                    expanded_batch_size = batch_size_local * num_passes
+                        [
+                            pass_spec
+                            for index, pass_spec in enumerate(pass_specs)
+                            if index % stage1_cfg_world_size == stage1_cfg_rank
+                        ]
+                        if stage1_cfg_parallel
+                        else pass_specs
+                    )
+                    num_execution_passes = len(execution_pass_specs)
+                    if num_execution_passes == 0:
+                        raise ValueError(
+                            "LTX2 stage-1 CFG parallel degree exceeds guidance pass count."
+                        )
+                    expanded_batch_size = batch_size_local * num_execution_passes
                     batched_model_kwargs = self._repeat_ltx2_model_kwargs_batch(
                         base_model_kwargs_local, expanded_batch_size
                     )
@@ -1733,21 +2101,21 @@ class LTX2DenoisingStage(DenoisingStage):
                         encoder_hidden_states=torch.cat(
                             [
                                 pass_spec.encoder_hidden_states
-                                for pass_spec in pass_specs
+                                for pass_spec in execution_pass_specs
                             ],
                             dim=0,
                         ),
                         audio_encoder_hidden_states=torch.cat(
                             [
                                 pass_spec.audio_encoder_hidden_states
-                                for pass_spec in pass_specs
+                                for pass_spec in execution_pass_specs
                             ],
                             dim=0,
                         ),
                         encoder_attention_mask=self._cat_or_none(
                             [
                                 pass_spec.encoder_attention_mask
-                                for pass_spec in pass_specs
+                                for pass_spec in execution_pass_specs
                             ]
                         ),
                     )
@@ -1755,14 +2123,14 @@ class LTX2DenoisingStage(DenoisingStage):
                         split_sizes = [1] * expanded_batch_size
                         split_pass_specs = tuple(
                             pass_spec
-                            for pass_spec in pass_specs
+                            for pass_spec in execution_pass_specs
                             for _ in range(batch_size_local)
                         )
                         split_perturbation_configs = (
                             ()
                             if use_split_pass_kwargs
                             else self._build_ltx2_guidance_perturbation_configs(
-                                pass_specs, batch_size_local
+                                execution_pass_specs, batch_size_local
                             )
                         )
                         batched_video_chunks = []
@@ -1796,7 +2164,7 @@ class LTX2DenoisingStage(DenoisingStage):
                     else:
                         perturbation_configs = (
                             self._build_ltx2_guidance_perturbation_configs(
-                                pass_specs, batch_size_local
+                                execution_pass_specs, batch_size_local
                             )
                         )
                         with self._ltx2_model_forward_context(ctx, step):
@@ -1813,18 +2181,20 @@ class LTX2DenoisingStage(DenoisingStage):
                             audio_chunk,
                         )
                         for pass_spec, video_chunk, audio_chunk in zip(
-                            pass_specs,
+                            execution_pass_specs,
-                            batched_video.chunk(num_passes, dim=0),
+                            batched_video.chunk(num_execution_passes, dim=0),
-                            batched_audio.chunk(num_passes, dim=0),
+                            batched_audio.chunk(num_execution_passes, dim=0),
                             strict=True,
                         )
                     }
+                    if not stage1_cfg_parallel:
                         v_pos, a_v_pos = pass_outputs["cond"]
                         v_neg, a_v_neg = pass_outputs["neg"]
                         v_ptb, a_v_ptb = pass_outputs.get("perturbed", (None, None))
                         v_mod, a_v_mod = pass_outputs.get("modality", (None, None))
 
                 sigma_value_float = float(sigma_value.item())
+                video_sigma_for_x0: float | torch.Tensor = sigma_value_float
                 audio_sigma_for_x0: float | torch.Tensor = sigma_value_float
                 if ctx.use_ltx23_hq_timestep_semantics:
                     video_sigma_for_x0 = model_inputs_local.timestep_video
@@ -1833,9 +2203,52 @@ class LTX2DenoisingStage(DenoisingStage):
                     video_sigma_for_x0 = sigma_value.to(
                         device=video_latents.device, dtype=torch.float32
                     ) * ctx.denoise_mask.squeeze(-1)
-                else:
-                    video_sigma_for_x0 = sigma_value_float
 
+                if stage1_cfg_parallel:
+                    guided_video = self._ltx2_combine_guided_x0_parallel(
+                        latents=video_latents,
+                        local_velocities={
+                            name: output[0] for name, output in pass_outputs.items()
+                        },
+                        sigma=video_sigma_for_x0,
+                        cfg_scale=float(stage1_guider_params["video_cfg_scale"]),
+                        stg_scale=float(stage1_guider_params["video_stg_scale"]),
+                        rescale_scale=float(
+                            stage1_guider_params["video_rescale_scale"]
+                        ),
+                        modality_scale=float(
+                            stage1_guider_params["video_modality_scale"]
+                        ),
+                    )
+                    if video_skip and ctx.last_denoised_video is not None:
+                        denoised_video_local = ctx.last_denoised_video
+                    else:
+                        denoised_video_local = guided_video
+                        if update_skip_cache:
+                            ctx.last_denoised_video = guided_video
+
+                    guided_audio = self._ltx2_combine_guided_x0_parallel(
+                        latents=audio_latents,
+                        local_velocities={
+                            name: output[1] for name, output in pass_outputs.items()
+                        },
+                        sigma=audio_sigma_for_x0,
+                        cfg_scale=float(stage1_guider_params["audio_cfg_scale"]),
+                        stg_scale=float(stage1_guider_params["audio_stg_scale"]),
+                        rescale_scale=float(
+                            stage1_guider_params["audio_rescale_scale"]
+                        ),
+                        modality_scale=float(
+                            stage1_guider_params["audio_modality_scale"]
+                        ),
+                    )
+                    if audio_skip and ctx.last_denoised_audio is not None:
+                        denoised_audio_local = ctx.last_denoised_audio
+                    else:
+                        denoised_audio_local = guided_audio
+                        if update_skip_cache:
+                            ctx.last_denoised_audio = guided_audio
+                else:
                     denoised_video_local = self._ltx2_velocity_to_x0(
                         video_latents, v_pos, video_sigma_for_x0
                     )
@@ -1892,8 +2305,12 @@ class LTX2DenoisingStage(DenoisingStage):
                         ),
                         cfg_scale=float(stage1_guider_params["video_cfg_scale"]),
                         stg_scale=float(stage1_guider_params["video_stg_scale"]),
-                    rescale_scale=float(stage1_guider_params["video_rescale_scale"]),
+                        rescale_scale=float(
-                    modality_scale=float(stage1_guider_params["video_modality_scale"]),
+                            stage1_guider_params["video_rescale_scale"]
+                        ),
+                        modality_scale=float(
+                            stage1_guider_params["video_modality_scale"]
+                        ),
                     )
                     if video_skip and ctx.last_denoised_video is not None:
                         denoised_video_local = ctx.last_denoised_video
@@ -1917,8 +2334,12 @@ class LTX2DenoisingStage(DenoisingStage):
                         ),
                         cfg_scale=float(stage1_guider_params["audio_cfg_scale"]),
                         stg_scale=float(stage1_guider_params["audio_stg_scale"]),
-                    rescale_scale=float(stage1_guider_params["audio_rescale_scale"]),
+                        rescale_scale=float(
-                    modality_scale=float(stage1_guider_params["audio_modality_scale"]),
+                            stage1_guider_params["audio_rescale_scale"]
+                        ),
+                        modality_scale=float(
+                            stage1_guider_params["audio_modality_scale"]
+                        ),
                     )
                     if audio_skip and ctx.last_denoised_audio is not None:
                         denoised_audio_local = ctx.last_denoised_audio
@@ -1968,14 +2389,22 @@ class LTX2DenoisingStage(DenoisingStage):
                 midpoint_audio_deterministic = anchor_audio + h * a21 * eps1_audio
 
                 substep_video_noise = (
-                    self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=True)
+                    self._ltx2_res2s_noise_like(
+                        ctx.latents, ctx, substep=True, batch=batch
+                    ).float()
                     if ctx.use_native_hq_res2s_sde_noise
                     else self._randn_like_with_batch_generators(
                         ctx.latents, batch
                     ).float()
                 )
                 substep_audio_noise = (
-                    self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=True)
+                    self._ltx2_res2s_noise_like(
+                        ctx.audio_latents,
+                        ctx,
+                        substep=True,
+                        batch=batch,
+                        is_audio=True,
+                    ).float()
                     if ctx.use_native_hq_res2s_sde_noise
                     else self._randn_like_with_batch_generators(
                         ctx.audio_latents, batch
@@ -2035,14 +2464,22 @@ class LTX2DenoisingStage(DenoisingStage):
                 )
 
                 step_video_noise = (
-                    self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=False)
+                    self._ltx2_res2s_noise_like(
+                        ctx.latents, ctx, substep=False, batch=batch
+                    ).float()
                     if ctx.use_native_hq_res2s_sde_noise
                     else self._randn_like_with_batch_generators(
                         ctx.latents, batch
                     ).float()
                 )
                 step_audio_noise = (
-                    self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=False)
+                    self._ltx2_res2s_noise_like(
+                        ctx.audio_latents,
+                        ctx,
+                        substep=False,
+                        batch=batch,
+                        is_audio=True,
+                    ).float()
                     if ctx.use_native_hq_res2s_sde_noise
                     else self._randn_like_with_batch_generators(
                         ctx.audio_latents, batch

python/sglang/multimodal_gen/runtime/server_args.py

@@ -160,6 +160,8 @@ class ServerArgs(DisaggArgsMixin):
     dp_degree: int = 1
     # cfg parallel (None = auto-decide based on num_gpus)
     enable_cfg_parallel: Optional[bool] = None
+    # number of GPUs in each CFG parallel group (None = auto, 1 = disabled, N > 1 = enabled)
+    cfg_parallel_degree: Optional[int] = None
 
     hsdp_replicate_dim: int = 1
     hsdp_shard_dim: Optional[int] = None
@@ -671,6 +673,14 @@ class ServerArgs(DisaggArgsMixin):
         if self.tp_size is None:
             self.tp_size = 1
 
+        # --cfg-parallel-size takes precedence over --enable-cfg-parallel bool.
+        if self.cfg_parallel_degree is not None:
+            if self.cfg_parallel_degree == 1:
+                self.enable_cfg_parallel = False
+            elif self.cfg_parallel_degree > 1:
+                self.enable_cfg_parallel = True
+            cfg_unspecified = False
+
         # Auto-enable CFG parallel when user hasn't set any parallelism flags
         # and there are enough GPUs.  Only auto-enable for models whose default
         # SamplingParams use classifier-free guidance (negative_prompt is not None),
@@ -695,11 +705,15 @@ class ServerArgs(DisaggArgsMixin):
             else:
                 self.enable_cfg_parallel = False
 
+        # Resolve cfg_parallel_degree to a concrete int now that enable_cfg_parallel is settled.
+        if self.cfg_parallel_degree is None:
+            self.cfg_parallel_degree = 2 if self.enable_cfg_parallel else 1
+
         # adjust sp_degree: allocate all remaining GPUs after TP and DP
         if self.sp_degree is None:
             num_gpus_per_group = self.dp_size * self.tp_size
             if self.enable_cfg_parallel:
-                num_gpus_per_group *= 2
+                num_gpus_per_group *= self.cfg_parallel_degree
             if self.num_gpus % num_gpus_per_group == 0:
                 self.sp_degree = self.num_gpus // num_gpus_per_group
             else:
@@ -739,10 +753,6 @@ class ServerArgs(DisaggArgsMixin):
             return False
         default_params = model_info.sampling_param_cls()
 
-        # for ltx2.3, cfg-parallel performs worse than ulysses-sp
-        is_ltx = "ltx" in type(default_params).__name__.lower()
-        if is_ltx:
-            return False
         return (
             getattr(default_params, "negative_prompt", None) is not None
             and getattr(default_params, "guidance_scale", 0) > 1.0
@@ -976,7 +986,18 @@ class ServerArgs(DisaggArgsMixin):
             "--enable-cfg-parallel",
             action="store_true",
             default=None,
-            help="Enable cfg parallel. Auto-enabled when num_gpus >= 2 and no SP flags are set.",
+            help="Enable cfg parallel at degree 2. Auto-enabled when num_gpus >= 2 and no SP flags are set.",
+        )
+        parser.add_argument(
+            "--cfg-parallel-size",
+            dest="cfg_parallel_degree",
+            type=int,
+            default=None,
+            help=(
+                "Number of GPUs per CFG parallel group (1 = disabled, N > 1 = enabled at degree N). "
+                "Supersedes --enable-cfg-parallel. Allows 4-branch CFG parallel (e.g., --cfg-parallel-size 4) "
+                "for models with cond + neg + perturbed + modality branches."
+            ),
         )
         parser.add_argument(
             "--data-parallel-size",
@@ -1627,11 +1648,13 @@ class ServerArgs(DisaggArgsMixin):
 
         num_gpus_per_group = self.dp_size * self.tp_size
         if self.enable_cfg_parallel:
-            num_gpus_per_group *= 2
+            num_gpus_per_group *= self.cfg_parallel_degree
 
         if self.num_gpus % num_gpus_per_group != 0:
             raise ValueError(
-                f"num_gpus ({self.num_gpus}) must be divisible by (dp_size * tp_size{' * 2' if self.enable_cfg_parallel else ''}) = {num_gpus_per_group}"
+                f"num_gpus ({self.num_gpus}) must be divisible by (dp_size * tp_size"
+                f"{f' * {self.cfg_parallel_degree}' if self.enable_cfg_parallel else ''}"
+                f") = {num_gpus_per_group}"
             )
 
         if self.sp_degree != self.ring_degree * self.ulysses_degree:

python/sglang/multimodal_gen/test/run_suite.py

@@ -86,7 +86,6 @@ STANDALONE_FILES = {
     "1-gpu": [
         "../cli/test_generate_t2i_perf.py",
         "test_update_weights_from_disk.py",
-        "test_tracing.py",
     ],
     "2-gpu": [
         "test_disagg_server.py",
@@ -100,7 +99,6 @@ STANDALONE_FILE_EST_TIMES = {
     "1-gpu": {
         "../cli/test_generate_t2i_perf.py": 240.0,
         "test_update_weights_from_disk.py": 480.0,
-        "test_tracing.py": 120.0,
     },
     "2-gpu": {
         # Two disagg clusters × (~3 min startup + ~1 min generate) ≈ 8 min.

python/sglang/multimodal_gen/test/server/accuracy_utils.py

@@ -14,6 +14,7 @@ from torch.distributed.tensor import distribute_tensor
 
 from sglang.multimodal_gen.runtime.distributed.parallel_state import (
     destroy_model_parallel,
+    get_classifier_free_guidance_world_size,
     get_data_parallel_world_size,
     get_sequence_parallel_world_size,
     get_tensor_model_parallel_world_size,
@@ -284,7 +285,7 @@ def initialize_parallel_runtime(sgl_args: ServerArgs) -> None:
     ulysses_degree = sgl_args.ulysses_degree
     ring_degree = sgl_args.ring_degree
     dp_size = sgl_args.dp_size
-    enable_cfg_parallel = bool(sgl_args.enable_cfg_parallel)
+    cfg_degree = sgl_args.cfg_parallel_degree or 1
 
     if (
         tp_size is None
@@ -305,7 +306,13 @@ def initialize_parallel_runtime(sgl_args: ServerArgs) -> None:
         current_tp = get_tensor_model_parallel_world_size()
         current_sp = get_sequence_parallel_world_size()
         current_dp = get_data_parallel_world_size()
-        if current_tp == tp_size and current_sp == sp_degree and current_dp == dp_size:
+        current_cfg = get_classifier_free_guidance_world_size()
+        if (
+            current_tp == tp_size
+            and current_sp == sp_degree
+            and current_dp == dp_size
+            and current_cfg == cfg_degree
+        ):
             return
         if torch.distributed.is_initialized():
             torch.distributed.barrier()
@@ -316,7 +323,7 @@ def initialize_parallel_runtime(sgl_args: ServerArgs) -> None:
     maybe_init_distributed_environment_and_model_parallel(
         tp_size=tp_size,
         sp_size=sp_degree,
-        enable_cfg_parallel=enable_cfg_parallel,
+        cfg_degree=cfg_degree,
         ulysses_degree=ulysses_degree,
         ring_degree=ring_degree,
         dp_size=dp_size,

python/sglang/multimodal_gen/test/server/gpu_cases.py

@@ -564,7 +564,7 @@ TWO_GPU_CASES = [
         "ltx_2_3_two_stage_ti2v_2gpus",
         DiffusionServerArgs(
             model_path="Lightricks/LTX-2.3",
-            ulysses_degree=2,
+            cfg_parallel=True,
             extras=[
                 "--pipeline-class-name LTX2TwoStagePipeline --ltx2-two-stage-device-mode original"
             ],
@@ -584,7 +584,7 @@ TWO_GPU_CASES = [
         "ltx_2.3_two_stage_t2v_2gpus",
         DiffusionServerArgs(
             model_path="Lightricks/LTX-2.3",
-            ulysses_degree=2,
+            cfg_parallel=True,
             extras=[
                 "--pipeline-class-name LTX2TwoStagePipeline",
                 "--ltx2-two-stage-device-mode original",

python/sglang/multimodal_gen/test/server/perf_baselines.json

@@ -2424,113 +2424,117 @@
         },
         "ltx_2.3_two_stage_t2v_2gpus": {
             "stages_ms": {
-                "InputValidationStage": 0.06,
+                "InputValidationStage": 0.07,
-                "TextEncodingStage": 1789.48,
+                "TextEncodingStage": 1784.13,
-                "LTX2TextConnectorStage": 28.01,
+                "LTX2TextConnectorStage": 29.83,
-                "LTX2HalveResolutionStage": 0.07,
+                "LTX2HalveResolutionStage": 0.06,
-                "LTX2LoRASwitchStage": 143.36,
+                "LTX2LoRASwitchStage": 153.74,
-                "LTX2SigmaPreparationStage": 0.37,
+                "LTX2SigmaPreparationStage": 0.63,
-                "TimestepPreparationStage": 24.63,
+                "TimestepPreparationStage": 31.17,
-                "LTX2AVLatentPreparationStage": 0.16,
+                "LTX2AVLatentPreparationStage": 0.26,
-                "LTX2AVDenoisingStage": 12652.72,
+                "LTX2ImageEncodingStage": 0.02,
-                "LTX2UpsampleStage": 779.33,
+                "LTX2AVDenoisingStage": 5460.05,
-                "LTX2RefinementStage": 762.78,
+                "LTX2UpsampleStage": 3.31,
-                "LTX2AVDecodingStage": 436.15
+                "LTX2RefinementStage": 615.01,
+                "LTX2AVDecodingStage": 240.26,
+                "per_frame_generation": null
             },
             "denoise_step_ms": {
-                "0": 4787.44,
+                "0": 179.26,
-                "1": 241.83,
+                "1": 283.15,
-                "2": 301.68,
+                "2": 193.3,
-                "3": 311.37,
+                "3": 166.82,
-                "4": 287.91,
+                "4": 200.16,
-                "5": 277.65,
+                "5": 196.51,
-                "6": 246.68,
+                "6": 178.52,
-                "7": 250.9,
+                "7": 168.41,
-                "8": 260.28,
+                "8": 182.67,
-                "9": 249.61,
+                "9": 166.79,
-                "10": 242.41,
+                "10": 163.28,
-                "11": 241.84,
+                "11": 170.97,
-                "12": 266.58,
+                "12": 176.57,
-                "13": 292.73,
+                "13": 189.92,
-                "14": 284.87,
+                "14": 162.76,
-                "15": 277.67,
+                "15": 179.89,
-                "16": 272.33,
+                "16": 172.41,
-                "17": 290.47,
+                "17": 184.39,
-                "18": 289.31,
+                "18": 161.55,
-                "19": 273.62,
+                "19": 176.63,
-                "20": 263.3,
+                "20": 200.48,
-                "21": 280.54,
+                "21": 188.38,
-                "22": 297.45,
+                "22": 161.98,
-                "23": 241.16,
+                "23": 172.25,
-                "24": 275.2,
+                "24": 167.45,
-                "25": 264.26,
+                "25": 189.81,
-                "26": 255.94,
+                "26": 176.08,
-                "27": 277.13,
+                "27": 170.88,
-                "28": 277.31,
+                "28": 188.96,
-                "29": 268.01,
+                "29": 183.21,
-                "30": 223.06,
+                "30": 238.92,
-                "31": 262.81,
+                "31": 187.04,
-                "32": 271.16
+                "32": 184.4
             },
-            "expected_e2e_ms": 22417.99,
+            "expected_e2e_ms": 18039.38,
-            "expected_avg_denoise_ms": 406.2,
+            "expected_avg_denoise_ms": 183.75,
-            "expected_median_denoise_ms": 272.33,
+            "expected_median_denoise_ms": 179.26,
-            "estimated_full_test_time_s": 216.7
+            "estimated_full_test_time_s": 160.0
         },
         "ltx_2_3_two_stage_ti2v_2gpus": {
             "stages_ms": {
-                "InputValidationStage": 4.02,
+                "InputValidationStage": 3.26,
-                "TextEncodingStage": 1681.03,
+                "TextEncodingStage": 1789.58,
-                "LTX2TextConnectorStage": 29.48,
+                "LTX2TextConnectorStage": 30.1,
-                "LTX2HalveResolutionStage": 0.06,
+                "LTX2HalveResolutionStage": 0.05,
-                "LTX2LoRASwitchStage": 123.34,
+                "LTX2LoRASwitchStage": 127.56,
-                "LTX2SigmaPreparationStage": 0.48,
+                "LTX2SigmaPreparationStage": 0.54,
-                "TimestepPreparationStage": 30.17,
+                "TimestepPreparationStage": 22.06,
-                "LTX2AVLatentPreparationStage": 0.23,
+                "LTX2AVLatentPreparationStage": 0.16,
-                "LTX2AVDenoisingStage": 28796.44,
+                "LTX2ImageEncodingStage": 27.81,
-                "LTX2UpsampleStage": 1062.28,
+                "LTX2AVDenoisingStage": 8517.54,
-                "LTX2RefinementStage": 1335.05,
+                "LTX2UpsampleStage": 2.6,
-                "LTX2AVDecodingStage": 482.26
+                "LTX2RefinementStage": 412.39,
+                "LTX2AVDecodingStage": 225.41,
+                "per_frame_generation": null
             },
             "denoise_step_ms": {
-                "0": 1205.82,
+                "0": 283.67,
-                "1": 1208.39,
+                "1": 332.51,
-                "2": 1269.93,
+                "2": 287.41,
-                "3": 998.3,
+                "3": 288.1,
-                "4": 921.32,
+                "4": 286.66,
-                "5": 1047.42,
+                "5": 286.86,
-                "6": 948.17,
+                "6": 285.56,
-                "7": 941.85,
+                "7": 285.01,
-                "8": 885.76,
+                "8": 281.79,
-                "9": 904.27,
+                "9": 276.41,
-                "10": 898.02,
+                "10": 279.14,
-                "11": 887.85,
+                "11": 306.21,
-                "12": 898.61,
+                "12": 280.42,
-                "13": 891.13,
+                "13": 279.25,
-                "14": 954.07,
+                "14": 282.39,
-                "15": 896.09,
+                "15": 283.28,
-                "16": 1001.81,
+                "16": 286.48,
-                "17": 881.83,
+                "17": 286.48,
-                "18": 882.7,
+                "18": 276.96,
-                "19": 920.33,
+                "19": 282.12,
-                "20": 971.85,
+                "20": 275.69,
-                "21": 891.13,
+                "21": 278.08,
-                "22": 884.23,
+                "22": 278.56,
-                "23": 870.44,
+                "23": 275.05,
-                "24": 890.55,
+                "24": 301.15,
-                "25": 869.53,
+                "25": 273.35,
-                "26": 861.77,
+                "26": 275.35,
-                "27": 873.08,
+                "27": 274.09,
-                "28": 934.75,
+                "28": 274.32,
-                "29": 865.63,
+                "29": 270.66,
-                "30": 357.92,
+                "30": 140.67,
-                "31": 315.56,
+                "31": 134.03,
-                "32": 345.67
+                "32": 132.51
             },
-            "expected_e2e_ms": 44740.03,
+            "expected_e2e_ms": 19149.77,
-            "expected_avg_denoise_ms": 890.17,
+            "expected_avg_denoise_ms": 270.31,
-            "expected_median_denoise_ms": 896.09,
+            "expected_median_denoise_ms": 280.42,
-            "estimated_full_test_time_s": 155.3
+            "estimated_full_test_time_s": 170.0
         },
         "ltx_2_3_hq_pipeline": {
             "stages_ms": {

python/sglang/multimodal_gen/test/server/test_disagg_server.py

@@ -543,10 +543,9 @@ class TestDisaggZImageTracing(_DisaggTestBase):
         # actual regression guard for this PR: it proves the W3C carrier
         # survives encoder→denoiser→decoder JSON hops (via ``_trace_state``).
         # The HTTP-level carrier extraction (root Req parented under the
-        # client's span_id) is already covered by ``test_tracing.py`` in
+        # client's span_id) is intentionally not asserted here: the server
-        # monolithic mode and asserting it here is flaky — the server head's
+        # head's BatchSpanProcessor may not flush the Req span before role
-        # BatchSpanProcessor may not flush the Req span before role spans
+        # spans reach the collector, since the role spans close first.
-        # reach the collector, since the role spans close first.
         trace_ids = {_as_hex(s.trace_id) for s in spans}
         self.assertEqual(
             trace_ids,

python/sglang/multimodal_gen/test/test_utils.py

@@ -33,7 +33,7 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
-SGL_TEST_FILES_CI_DATA_REVISION = "6e7b99e16b857c98285277fe3b4ffef30559bde9"
+SGL_TEST_FILES_CI_DATA_REVISION = "3ca3bad088ecc9ef80947d85c551cd335c75b87f"
 SGL_TEST_FILES_CONSISTENCY_GT_ROOT = (
     "https://raw.githubusercontent.com/"
     f"sgl-project/ci-data/{SGL_TEST_FILES_CI_DATA_REVISION}/"

python/sglang/multimodal_gen/test/unit/test_cfg_policy.py

@@ -0,0 +1,33 @@
+import unittest
+from unittest.mock import MagicMock
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed.cfg_policy import CFGPolicy
+
+
+class TestCFGPolicyCombine(unittest.TestCase):
+    def test_cfg_parallel_uses_parallel_arithmetic_order(self):
+        policy = CFGPolicy()
+        req = MagicMock()
+        req.cfg_normalization = 0
+        req.guidance_rescale = 0
+
+        pipeline_config = MagicMock()
+        pipeline_config.postprocess_cfg_noise.side_effect = lambda _, noise, __: noise
+
+        pos = torch.tensor([1.0], dtype=torch.bfloat16)
+        neg = torch.tensor([0.1], dtype=torch.bfloat16)
+
+        serial = policy.combine([pos, neg], req, 7.0, pipeline_config)
+        parallel = policy.combine(
+            [pos, neg], req, 7.0, pipeline_config, cfg_parallel=True
+        )
+
+        self.assertTrue(torch.equal(serial, neg + 7.0 * (pos - neg)))
+        self.assertTrue(torch.equal(parallel, 7.0 * pos + (1 - 7.0) * neg))
+        self.assertFalse(torch.equal(serial, parallel))
+
+
+if __name__ == "__main__":
+    unittest.main()