From 9c41b1058fbe4ae911c4faeedbc7efcea88c1c4c Mon Sep 17 00:00:00 2001 From: ykcai-daniel Date: Thu, 7 May 2026 07:56:55 -0700 Subject: [PATCH] [diffusion] refactor: refactor cfg parallelism framework to support multi-branch CFG for LTX2 (#23736) Co-authored-by: Mick --- .../scripts/bench_diffusion_denoise.py | 22 +- .../configs/pipeline_configs/base.py | 2 + .../runtime/distributed/cfg_parallel_utils.py | 181 +++++ .../runtime/distributed/cfg_policy.py | 159 ++++ .../runtime/distributed/parallel_state.py | 4 +- .../runtime/managers/gpu_worker.py | 2 +- .../pipelines_core/stages/denoising.py | 369 ++------- .../pipelines_core/stages/denoising_av.py | 11 + .../pipelines_core/stages/ltx_2_denoising.py | 739 ++++++++++++++---- .../multimodal_gen/runtime/server_args.py | 39 +- .../sglang/multimodal_gen/test/run_suite.py | 2 - .../test/server/accuracy_utils.py | 13 +- .../multimodal_gen/test/server/gpu_cases.py | 4 +- .../test/server/perf_baselines.json | 200 ++--- .../test/server/test_disagg_server.py | 7 +- .../sglang/multimodal_gen/test/test_utils.py | 2 +- .../test/unit/test_cfg_policy.py | 33 + 17 files changed, 1225 insertions(+), 564 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/distributed/cfg_parallel_utils.py create mode 100644 python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_cfg_policy.py diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py index 2e93b2c8e..4113bc419 100755 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py +++ b/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: /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", diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index e3be31100..7005ac70d 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/distributed/cfg_parallel_utils.py b/python/sglang/multimodal_gen/runtime/distributed/cfg_parallel_utils.py new file mode 100644 index 000000000..39a289516 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py b/python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py new file mode 100644 index 000000000..b65e69fde --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index cc2c7d443..e4bff7ed6 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 02438fd07..1f282b221 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 2901b50c7..ef56c3b8a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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_rank, + get_classifier_free_guidance_world_size, ) 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, - pos_cond_kwargs=ctx.pos_cond_kwargs, - neg_cond_kwargs=ctx.neg_cond_kwargs, + cfg_policy=ctx.cfg_policy, 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( - 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( + def _predict_noise_with_cfg( self, + current_model: nn.Module, + latent_model_input: torch.Tensor, + timestep, batch: Req, - noise_pred: torch.Tensor, - noise_pred_cond: torch.Tensor | None, - cfg_rank: int, - ) -> torch.Tensor: - # keep model-specific CFG behavior out of the main denoising loop - # for cfg-parallel, broadcast cond noise first so the hook sees the same - # inputs as the serial path. - if cfg_rank == 0: - assert noise_pred_cond is not None - cond_noise = noise_pred_cond + timestep_index: int, + attn_metadata, + target_dtype, + current_guidance_scale, + cfg_policy: CFGPolicy, + server_args: ServerArgs, + guidance: torch.Tensor, + latents: torch.Tensor, + ) -> "torch.Tensor | tuple[torch.Tensor, ...]": + """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? - cond_noise = torch.empty_like(noise_pred) - cond_noise = get_cfg_group().broadcast(cond_noise, src=0) + # perform cfg branches one-by-one locally + predictions = [predict_fn(branch) for branch in cfg_policy.branches] - # qwen-image uses true_cfg_scale, match the per-token norm back to the conditional branch - return self.server_args.pipeline_config.postprocess_cfg_noise( - batch, noise_pred, cond_noise - ) - - def _combine_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 + return cfg_policy.combine( + predictions, + batch, + cfg_scale, + server_args.pipeline_config, + cfg_parallel=server_args.enable_cfg_parallel, ) 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. diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py index df26a38b5..6fb7a0280 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py index c6a0fb492..2328558c2 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py +++ b/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,53 +1655,90 @@ class LTX2DenoisingStage(DenoisingStage): ) use_official_cfg_path = stage1_guider_params is None if use_official_cfg_path: - model_kwargs = self._build_ltx2_model_kwargs( - ctx, - base_model_kwargs, - encoder_hidden_states=batch.prompt_embeds[0], - audio_encoder_hidden_states=batch.audio_prompt_embeds[0], - encoder_attention_mask=prompt_attention_mask, + cfg_parallel = ( + server_args.enable_cfg_parallel and batch.do_classifier_free_guidance ) - if batch.do_classifier_free_guidance: - cfg_batch_size = batch_size * 2 - model_kwargs = self._repeat_ltx2_model_kwargs_batch( - model_kwargs, cfg_batch_size + 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, ) - model_kwargs["encoder_hidden_states"] = torch.cat( - [batch.negative_prompt_embeds[0], batch.prompt_embeds[0]], dim=0 + else: + model_kwargs = self._build_ltx2_model_kwargs( + ctx, + base_model_kwargs, + encoder_hidden_states=batch.prompt_embeds[0], + audio_encoder_hidden_states=batch.audio_prompt_embeds[0], + encoder_attention_mask=prompt_attention_mask, ) - model_kwargs["audio_encoder_hidden_states"] = torch.cat( - [ - batch.negative_audio_prompt_embeds[0], - batch.audio_prompt_embeds[0], - ], - dim=0, - ) - if self._should_pass_ltx2_text_attention_mask(ctx): - repeated_attention_mask = self._cat_or_none( + if batch.do_classifier_free_guidance: + cfg_batch_size = batch_size * 2 + model_kwargs = self._repeat_ltx2_model_kwargs_batch( + model_kwargs, cfg_batch_size + ) + model_kwargs["encoder_hidden_states"] = torch.cat( + [batch.negative_prompt_embeds[0], batch.prompt_embeds[0]], + dim=0, + ) + model_kwargs["audio_encoder_hidden_states"] = torch.cat( [ - self._get_ltx_prompt_attention_mask( - batch, - is_ltx23_variant=( - ctx.is_ltx23_variant - and not ctx.use_ltx23_legacy_one_stage + batch.negative_audio_prompt_embeds[0], + batch.audio_prompt_embeds[0], + ], + dim=0, + ) + if self._should_pass_ltx2_text_attention_mask(ctx): + repeated_attention_mask = self._cat_or_none( + [ + self._get_ltx_prompt_attention_mask( + batch, + is_ltx23_variant=( + ctx.is_ltx23_variant + and not ctx.use_ltx23_legacy_one_stage + ), + negative=True, ), - negative=True, - ), - prompt_attention_mask, - ] - ) - model_kwargs["encoder_attention_mask"] = repeated_attention_mask - model_kwargs["audio_encoder_attention_mask"] = ( - repeated_attention_mask - ) + prompt_attention_mask, + ] + ) + model_kwargs["encoder_attention_mask"] = repeated_attention_mask + model_kwargs["audio_encoder_attention_mask"] = ( + repeated_attention_mask + ) with self._ltx2_model_forward_context(ctx, step): model_video, model_audio = step.current_model(**model_kwargs) 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) - expanded_batch_size = batch_size_local * num_passes + execution_pass_specs = ( + [ + 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, - batched_video.chunk(num_passes, dim=0), - batched_audio.chunk(num_passes, dim=0), + execution_pass_specs, + batched_video.chunk(num_execution_passes, dim=0), + batched_audio.chunk(num_execution_passes, dim=0), strict=True, ) } - 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)) + 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,99 +2203,150 @@ 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 - denoised_video_local = self._ltx2_velocity_to_x0( - video_latents, v_pos, video_sigma_for_x0 - ) - denoised_audio_local = self._ltx2_velocity_to_x0( - audio_latents, a_v_pos, audio_sigma_for_x0 - ) - denoised_video_neg = self._ltx2_velocity_to_x0( - video_latents, v_neg, video_sigma_for_x0 - ) - denoised_audio_neg = self._ltx2_velocity_to_x0( - audio_latents, a_v_neg, audio_sigma_for_x0 - ) - denoised_video_perturbed = ( - None - if v_ptb is None - else self._ltx2_velocity_to_x0( - video_latents, v_ptb, video_sigma_for_x0 + 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"] + ), ) - ) - denoised_audio_perturbed = ( - None - if a_v_ptb is None - else self._ltx2_velocity_to_x0( - audio_latents, a_v_ptb, audio_sigma_for_x0 - ) - ) - denoised_video_modality = ( - None - if v_mod is None - else self._ltx2_velocity_to_x0( - video_latents, v_mod, video_sigma_for_x0 - ) - ) - denoised_audio_modality = ( - None - if a_v_mod is None - else self._ltx2_velocity_to_x0( - audio_latents, a_v_mod, audio_sigma_for_x0 - ) - ) + 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_video = self._ltx2_calculate_guided_x0( - cond=denoised_video_local, - uncond_text=denoised_video_neg, - uncond_perturbed=( - denoised_video_perturbed - if denoised_video_perturbed is not None - else 0.0 - ), - uncond_modality=( - denoised_video_modality - if denoised_video_modality is not None - else 0.0 - ), - 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 + 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 = guided_video - if update_skip_cache: - ctx.last_denoised_video = guided_video + denoised_video_local = self._ltx2_velocity_to_x0( + video_latents, v_pos, video_sigma_for_x0 + ) + denoised_audio_local = self._ltx2_velocity_to_x0( + audio_latents, a_v_pos, audio_sigma_for_x0 + ) + denoised_video_neg = self._ltx2_velocity_to_x0( + video_latents, v_neg, video_sigma_for_x0 + ) + denoised_audio_neg = self._ltx2_velocity_to_x0( + audio_latents, a_v_neg, audio_sigma_for_x0 + ) + denoised_video_perturbed = ( + None + if v_ptb is None + else self._ltx2_velocity_to_x0( + video_latents, v_ptb, video_sigma_for_x0 + ) + ) + denoised_audio_perturbed = ( + None + if a_v_ptb is None + else self._ltx2_velocity_to_x0( + audio_latents, a_v_ptb, audio_sigma_for_x0 + ) + ) + denoised_video_modality = ( + None + if v_mod is None + else self._ltx2_velocity_to_x0( + video_latents, v_mod, video_sigma_for_x0 + ) + ) + denoised_audio_modality = ( + None + if a_v_mod is None + else self._ltx2_velocity_to_x0( + audio_latents, a_v_mod, audio_sigma_for_x0 + ) + ) - guided_audio = self._ltx2_calculate_guided_x0( - cond=denoised_audio_local, - uncond_text=denoised_audio_neg, - uncond_perturbed=( - denoised_audio_perturbed - if denoised_audio_perturbed is not None - else 0.0 - ), - uncond_modality=( - denoised_audio_modality - if denoised_audio_modality is not None - else 0.0 - ), - 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 + guided_video = self._ltx2_calculate_guided_x0( + cond=denoised_video_local, + uncond_text=denoised_video_neg, + uncond_perturbed=( + denoised_video_perturbed + if denoised_video_perturbed is not None + else 0.0 + ), + uncond_modality=( + denoised_video_modality + if denoised_video_modality is not None + else 0.0 + ), + 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_calculate_guided_x0( + cond=denoised_audio_local, + uncond_text=denoised_audio_neg, + uncond_perturbed=( + denoised_audio_perturbed + if denoised_audio_perturbed is not None + else 0.0 + ), + uncond_modality=( + denoised_audio_modality + if denoised_audio_modality is not None + else 0.0 + ), + 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 denoised_video_local = self._ltx2_apply_clean_latent_mask( denoised_video_local, ctx @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 6e31b35b9..5e1b02e20 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/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: diff --git a/python/sglang/multimodal_gen/test/run_suite.py b/python/sglang/multimodal_gen/test/run_suite.py index 2ddfa9975..1c7c41be4 100644 --- a/python/sglang/multimodal_gen/test/run_suite.py +++ b/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. diff --git a/python/sglang/multimodal_gen/test/server/accuracy_utils.py b/python/sglang/multimodal_gen/test/server/accuracy_utils.py index 5d6dcf385..cd544e6b9 100644 --- a/python/sglang/multimodal_gen/test/server/accuracy_utils.py +++ b/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, diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index 8abc50665..9407625a5 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/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", diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 648b92510..08685b75c 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/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, - "TextEncodingStage": 1789.48, - "LTX2TextConnectorStage": 28.01, - "LTX2HalveResolutionStage": 0.07, - "LTX2LoRASwitchStage": 143.36, - "LTX2SigmaPreparationStage": 0.37, - "TimestepPreparationStage": 24.63, - "LTX2AVLatentPreparationStage": 0.16, - "LTX2AVDenoisingStage": 12652.72, - "LTX2UpsampleStage": 779.33, - "LTX2RefinementStage": 762.78, - "LTX2AVDecodingStage": 436.15 + "InputValidationStage": 0.07, + "TextEncodingStage": 1784.13, + "LTX2TextConnectorStage": 29.83, + "LTX2HalveResolutionStage": 0.06, + "LTX2LoRASwitchStage": 153.74, + "LTX2SigmaPreparationStage": 0.63, + "TimestepPreparationStage": 31.17, + "LTX2AVLatentPreparationStage": 0.26, + "LTX2ImageEncodingStage": 0.02, + "LTX2AVDenoisingStage": 5460.05, + "LTX2UpsampleStage": 3.31, + "LTX2RefinementStage": 615.01, + "LTX2AVDecodingStage": 240.26, + "per_frame_generation": null }, "denoise_step_ms": { - "0": 4787.44, - "1": 241.83, - "2": 301.68, - "3": 311.37, - "4": 287.91, - "5": 277.65, - "6": 246.68, - "7": 250.9, - "8": 260.28, - "9": 249.61, - "10": 242.41, - "11": 241.84, - "12": 266.58, - "13": 292.73, - "14": 284.87, - "15": 277.67, - "16": 272.33, - "17": 290.47, - "18": 289.31, - "19": 273.62, - "20": 263.3, - "21": 280.54, - "22": 297.45, - "23": 241.16, - "24": 275.2, - "25": 264.26, - "26": 255.94, - "27": 277.13, - "28": 277.31, - "29": 268.01, - "30": 223.06, - "31": 262.81, - "32": 271.16 + "0": 179.26, + "1": 283.15, + "2": 193.3, + "3": 166.82, + "4": 200.16, + "5": 196.51, + "6": 178.52, + "7": 168.41, + "8": 182.67, + "9": 166.79, + "10": 163.28, + "11": 170.97, + "12": 176.57, + "13": 189.92, + "14": 162.76, + "15": 179.89, + "16": 172.41, + "17": 184.39, + "18": 161.55, + "19": 176.63, + "20": 200.48, + "21": 188.38, + "22": 161.98, + "23": 172.25, + "24": 167.45, + "25": 189.81, + "26": 176.08, + "27": 170.88, + "28": 188.96, + "29": 183.21, + "30": 238.92, + "31": 187.04, + "32": 184.4 }, - "expected_e2e_ms": 22417.99, - "expected_avg_denoise_ms": 406.2, - "expected_median_denoise_ms": 272.33, - "estimated_full_test_time_s": 216.7 + "expected_e2e_ms": 18039.38, + "expected_avg_denoise_ms": 183.75, + "expected_median_denoise_ms": 179.26, + "estimated_full_test_time_s": 160.0 }, "ltx_2_3_two_stage_ti2v_2gpus": { "stages_ms": { - "InputValidationStage": 4.02, - "TextEncodingStage": 1681.03, - "LTX2TextConnectorStage": 29.48, - "LTX2HalveResolutionStage": 0.06, - "LTX2LoRASwitchStage": 123.34, - "LTX2SigmaPreparationStage": 0.48, - "TimestepPreparationStage": 30.17, - "LTX2AVLatentPreparationStage": 0.23, - "LTX2AVDenoisingStage": 28796.44, - "LTX2UpsampleStage": 1062.28, - "LTX2RefinementStage": 1335.05, - "LTX2AVDecodingStage": 482.26 + "InputValidationStage": 3.26, + "TextEncodingStage": 1789.58, + "LTX2TextConnectorStage": 30.1, + "LTX2HalveResolutionStage": 0.05, + "LTX2LoRASwitchStage": 127.56, + "LTX2SigmaPreparationStage": 0.54, + "TimestepPreparationStage": 22.06, + "LTX2AVLatentPreparationStage": 0.16, + "LTX2ImageEncodingStage": 27.81, + "LTX2AVDenoisingStage": 8517.54, + "LTX2UpsampleStage": 2.6, + "LTX2RefinementStage": 412.39, + "LTX2AVDecodingStage": 225.41, + "per_frame_generation": null }, "denoise_step_ms": { - "0": 1205.82, - "1": 1208.39, - "2": 1269.93, - "3": 998.3, - "4": 921.32, - "5": 1047.42, - "6": 948.17, - "7": 941.85, - "8": 885.76, - "9": 904.27, - "10": 898.02, - "11": 887.85, - "12": 898.61, - "13": 891.13, - "14": 954.07, - "15": 896.09, - "16": 1001.81, - "17": 881.83, - "18": 882.7, - "19": 920.33, - "20": 971.85, - "21": 891.13, - "22": 884.23, - "23": 870.44, - "24": 890.55, - "25": 869.53, - "26": 861.77, - "27": 873.08, - "28": 934.75, - "29": 865.63, - "30": 357.92, - "31": 315.56, - "32": 345.67 + "0": 283.67, + "1": 332.51, + "2": 287.41, + "3": 288.1, + "4": 286.66, + "5": 286.86, + "6": 285.56, + "7": 285.01, + "8": 281.79, + "9": 276.41, + "10": 279.14, + "11": 306.21, + "12": 280.42, + "13": 279.25, + "14": 282.39, + "15": 283.28, + "16": 286.48, + "17": 286.48, + "18": 276.96, + "19": 282.12, + "20": 275.69, + "21": 278.08, + "22": 278.56, + "23": 275.05, + "24": 301.15, + "25": 273.35, + "26": 275.35, + "27": 274.09, + "28": 274.32, + "29": 270.66, + "30": 140.67, + "31": 134.03, + "32": 132.51 }, - "expected_e2e_ms": 44740.03, - "expected_avg_denoise_ms": 890.17, - "expected_median_denoise_ms": 896.09, - "estimated_full_test_time_s": 155.3 + "expected_e2e_ms": 19149.77, + "expected_avg_denoise_ms": 270.31, + "expected_median_denoise_ms": 280.42, + "estimated_full_test_time_s": 170.0 }, "ltx_2_3_hq_pipeline": { "stages_ms": { diff --git a/python/sglang/multimodal_gen/test/server/test_disagg_server.py b/python/sglang/multimodal_gen/test/server/test_disagg_server.py index a8233a6fd..c3aee6ff3 100755 --- a/python/sglang/multimodal_gen/test/server/test_disagg_server.py +++ b/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 - # monolithic mode and asserting it here is flaky — the server head's - # BatchSpanProcessor may not flush the Req span before role spans - # reach the collector, since the role spans close first. + # client's span_id) is intentionally not asserted here: the server + # head's BatchSpanProcessor may not flush the Req span before role + # spans reach the collector, since the role spans close first. trace_ids = {_as_hex(s.trace_id) for s in spans} self.assertEqual( trace_ids, diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 5e96bbcb0..ccda47743 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/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}/" diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_policy.py b/python/sglang/multimodal_gen/test/unit/test_cfg_policy.py new file mode 100644 index 000000000..09d48eaad --- /dev/null +++ b/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()