[diffusion] optimize: optimize LTX2.3 CFG path (#28624)

python/sglang/multimodal_gen/configs/models/dits/ltx_2.py

@@ -162,6 +162,7 @@ class LTX2ArchConfig(DiTArchConfig):
     # SGLang-specific parameters
     patch_size: tuple[int, int, int] = (1, 2, 2)
     text_len: int = 512
+    enable_packed_qkv_input_a2a: bool = False
 
     def __post_init__(self):
         super().__post_init__()

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

@@ -202,6 +202,7 @@ class LTX2PipelineConfig(PipelineConfig):
         return ModelDeploymentConfig(
             auto_disable_component_offload_min_available_memory_gb=70,
             auto_disable_component_offload_components=("dit",),
+            auto_cfg_parallel_degree_by_num_gpus=((4, 1), (8, 1)),
         )
 
     def prepare_latent_shape(self, batch, batch_size, num_frames):

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

@@ -25,3 +25,11 @@ class ModelDeploymentConfig:
     fsdp_auto_requires_cfg: bool = True
     fsdp_auto_requires_default_parallelism: bool = True
     auto_enable_cfg_parallel: bool = True
+    # degree 1 keeps CFG parallel disabled and leaves GPUs available for SP
+    auto_cfg_parallel_degree_by_num_gpus: tuple[tuple[int, int], ...] = ()
+
+    def get_auto_cfg_parallel_degree(self, num_gpus: int) -> int:
+        for candidate_num_gpus, cfg_degree in self.auto_cfg_parallel_degree_by_num_gpus:
+            if candidate_num_gpus == num_gpus:
+                return cfg_degree
+        return 2

python/sglang/multimodal_gen/runtime/layers/attention/layer.py

@@ -425,6 +425,7 @@ class USPAttention(nn.Module):
         prefix: str = "",
         dropout_rate: float = 0.0,
         skip_sequence_parallel: bool = False,
+        enable_packed_qkv_input_a2a: bool = False,
         **extra_impl_args,
     ) -> None:
         """
@@ -480,6 +481,7 @@ class USPAttention(nn.Module):
         self.dropout_p = dropout_rate
 
         self.skip_sequence_parallel = skip_sequence_parallel
+        self.enable_packed_qkv_input_a2a = bool(enable_packed_qkv_input_a2a)
 
     def _get_usp_a2a_stream(self):
         if USPAttention._usp_a2a_stream is None:
@@ -781,6 +783,18 @@ class USPAttention(nn.Module):
         # Ulysses-style All-to-All for sequence/head sharding
         if sp_size > 1:
             # -> [B, S, H_local, D]
+            if self.enable_packed_qkv_input_a2a and q.device.type == "cuda":
+                q, k, v = async_a2a_communicate(
+                    [q, k, v],
+                    sp_size,
+                    get_sp_group().ulysses_group,
+                    self._get_usp_a2a_stream(),
+                    local_seq_2_local_head=True,
+                )
+                q = q.contiguous()
+                k = k.contiguous()
+                v = v.contiguous()
+            else:
                 q = _usp_input_all_to_all(q, head_dim=2)
                 k = _usp_input_all_to_all(k, head_dim=2)
                 v = _usp_input_all_to_all(v, head_dim=2)

python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py

@@ -604,6 +604,7 @@ class LTX2Attention(nn.Module):
         qk_norm: bool = True,
         use_local_attention: bool = False,
         apply_gated_attention: bool = False,
+        enable_packed_qkv_input_a2a: bool = False,
         supported_attention_backends: set[AttentionBackendEnum] | None = None,
         prefix: str = "",
         quant_config: QuantizationConfig | None = None,
@@ -619,6 +620,7 @@ class LTX2Attention(nn.Module):
         self.qk_norm = bool(qk_norm)
         self.use_local_attention = bool(use_local_attention)
         self.apply_gated_attention = bool(apply_gated_attention)
+        self.enable_packed_qkv_input_a2a = bool(enable_packed_qkv_input_a2a)
         self.prefix = prefix
 
         tp_size = get_tp_world_size()
@@ -706,6 +708,7 @@ class LTX2Attention(nn.Module):
                 causal=False,
                 supported_attention_backends=supported_attention_backends,
                 prefix=f"{prefix}.attn",
+                enable_packed_qkv_input_a2a=self.enable_packed_qkv_input_a2a,
                 # official LTX2 torch_sdpa uses cuDNN; cuda setup disables it
                 allow_cudnn_sdp=True,
             )
@@ -931,6 +934,7 @@ class LTX2TransformerBlock(nn.Module):
         cross_attention_adaln: bool = False,
         use_local_av_cross_attention: bool = False,
         force_sdpa_v2a_cross_attention: bool = False,
+        enable_packed_qkv_input_a2a: bool = False,
         supported_attention_backends: set[AttentionBackendEnum] | None = None,
         prefix: str = "",
         quant_config: QuantizationConfig | None = None,
@@ -951,6 +955,7 @@ class LTX2TransformerBlock(nn.Module):
             norm_eps=norm_eps,
             qk_norm=qk_norm,
             apply_gated_attention=apply_gated_attention,
+            enable_packed_qkv_input_a2a=enable_packed_qkv_input_a2a,
             supported_attention_backends=supported_attention_backends,
             prefix=f"{prefix}.attn1",
             quant_config=quant_config,
@@ -962,6 +967,7 @@ class LTX2TransformerBlock(nn.Module):
             norm_eps=norm_eps,
             qk_norm=qk_norm,
             apply_gated_attention=apply_gated_attention,
+            enable_packed_qkv_input_a2a=enable_packed_qkv_input_a2a,
             supported_attention_backends=supported_attention_backends,
             prefix=f"{prefix}.audio_attn1",
             quant_config=quant_config,
@@ -1007,6 +1013,7 @@ class LTX2TransformerBlock(nn.Module):
             qk_norm=qk_norm,
             use_local_attention=use_local_av_cross_attention,
             apply_gated_attention=apply_gated_attention,
+            enable_packed_qkv_input_a2a=enable_packed_qkv_input_a2a,
             supported_attention_backends=supported_attention_backends,
             prefix=f"{prefix}.audio_to_video_attn",
             quant_config=quant_config,
@@ -1020,6 +1027,7 @@ class LTX2TransformerBlock(nn.Module):
             qk_norm=qk_norm,
             use_local_attention=use_local_av_cross_attention,
             apply_gated_attention=apply_gated_attention,
+            enable_packed_qkv_input_a2a=enable_packed_qkv_input_a2a,
             supported_attention_backends=(
                 {AttentionBackendEnum.TORCH_SDPA}
                 if force_sdpa_v2a_cross_attention
@@ -1665,6 +1673,7 @@ class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
                     force_sdpa_v2a_cross_attention=bool(
                         getattr(arch, "force_sdpa_v2a_cross_attention", False)
                     ),
+                    enable_packed_qkv_input_a2a=arch.enable_packed_qkv_input_a2a,
                     supported_attention_backends=self._supported_attention_backends,
                     prefix=config.prefix,
                     quant_config=quant_config,

python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py

@@ -68,7 +68,22 @@ class ParallelExecutor(PipelineExecutor):
                     torch.distributed.barrier()
 
                 elif paradigm == StageParallelismType.CFG_PARALLEL:
+                    local_batch = batch
+                    local_batch_fields = stage.cfg_parallel_local_batch_fields(
+                        batch, server_args
+                    )
+                    # filter local batch fields from batch
+                    if rank == 0 and local_batch_fields:
+                        local_field_values = {
+                            name: getattr(batch, name) for name in local_batch_fields
+                        }
+                        for name in local_batch_fields:
+                            setattr(batch, name, None)
+                    else:
+                        local_field_values = {}
+
                     obj_list = [batch] if rank == 0 else []
+                    try:
                         # `dist.broadcast(src=...)` expects a global rank for process groups.
                         broadcasted_list = broadcast_pyobj(
                             obj_list,
@@ -76,8 +91,15 @@ class ParallelExecutor(PipelineExecutor):
                             dist_group=cfg_group.cpu_group,
                             src=cfg_group.ranks[0],
                         )
+                    finally:
+                        if rank == 0:
+                            # resume local batch fields on rank 0
+                            for name, value in local_field_values.items():
+                                setattr(batch, name, value)
                     if rank != 0:
                         batch = broadcasted_list[0]
+                        for name in local_batch_fields:
+                            setattr(batch, name, getattr(local_batch, name))
                     batch = self._run_stage_with_executor_hooks(
                         stage,
                         stage_index,

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

@@ -272,6 +272,12 @@ class PipelineStage(StageDedupMixin, ABC):
         #     return StageParallelismType.MAIN_RANK_ONLY
         return StageParallelismType.REPLICATED
 
+    def cfg_parallel_local_batch_fields(
+        self, batch: Req, server_args: ServerArgs
+    ) -> tuple[str, ...]:
+        """the name of fields which already have a local version on each GPU in CFG-Parallel, no need to broadcast"""
+        return ()
+
     def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
         """
         Verify the output for the stage.

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py

@@ -164,10 +164,13 @@ class LTX2AVDecodingStage(DecodingStage):
                     audio_vae_dtype, server_args.disable_autocast
                 )
                 should_cast_audio_vae = not audio_vae_autocast_enabled
-                with torch.no_grad(), torch.autocast(
+                with (
+                    torch.no_grad(),
+                    torch.autocast(
                         device_type=current_platform.device_type,
                         dtype=audio_vae_dtype,
                         enabled=audio_vae_autocast_enabled,
+                    ),
                 ):
                     # Decode latents to spectrogram
                     with temporary_module_dtype(

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

@@ -167,6 +167,13 @@ class LTX2DenoisingStage(DenoisingStage):
             return StageParallelismType.CFG_PARALLEL
         return StageParallelismType.REPLICATED
 
+    def cfg_parallel_local_batch_fields(
+        self, batch: Req, server_args: ServerArgs
+    ) -> tuple[str, ...]:
+        if is_ltx23_native_variant(server_args.pipeline_config.vae_config.arch_config):
+            return ("latents", "audio_latents")
+        return ()
+
     @staticmethod
     def _combine_cfg_parallel_av(
         video: torch.Tensor,

python/sglang/multimodal_gen/runtime/server_args.py

@@ -821,7 +821,13 @@ class ServerArgs(DisaggServerArgsMixin):
         # because non-CFG models (e.g. FLUX) crash when CFG parallel splits ranks.
         if cfg_unspecified:
             deployment_config = self.pipeline_config.get_model_deployment_config()
-            cfg_group_size = self.dp_size * self.tp_size * 2
+            auto_cfg_parallel_degree = deployment_config.get_auto_cfg_parallel_degree(
+                self.num_gpus
+            )
+            if auto_cfg_parallel_degree < 1:
+                self.enable_cfg_parallel = False
+            else:
+                cfg_group_size = self.dp_size * self.tp_size * auto_cfg_parallel_degree
                 if (
                     self.performance_mode != "manual"
                     and deployment_config.auto_enable_cfg_parallel
@@ -832,11 +838,19 @@ class ServerArgs(DisaggServerArgsMixin):
                     and ring_unspecified
                     and self._model_default_uses_cfg()
                 ):
-                self.enable_cfg_parallel = True
+                    self.cfg_parallel_degree = auto_cfg_parallel_degree
+                    self.enable_cfg_parallel = auto_cfg_parallel_degree > 1
+                    if self.enable_cfg_parallel:
                         logger.info(
-                    "Automatically enabled CFG parallel for %d GPUs. "
+                            "Automatically enabled CFG parallel at degree %d for %d GPUs. "
                             "Use --sp-degree / --ulysses-degree to use sequence "
                             "parallelism instead.",
+                            self.cfg_parallel_degree,
+                            self.num_gpus,
+                        )
+                    else:
+                        logger.info(
+                            "Automatically disabled CFG parallel for %d GPUs based on model deployment config.",
                             self.num_gpus,
                         )
                 else:

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

@@ -15,7 +15,10 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ModelTaskType,
     PipelineConfig,
 )
-from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig
+from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
+    LTX2PipelineConfig,
+    LTX23PipelineConfig,
+)
 from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig
 from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
     QwenImagePipelineConfig,
@@ -795,6 +798,7 @@ class TestOffloadDefaults(unittest.TestCase):
         mova_deployment = MOVAPipelineConfig().get_model_deployment_config()
         zimage_deployment = ZImagePipelineConfig().get_model_deployment_config()
         ltx_deployment = LTX2PipelineConfig().get_model_deployment_config()
+        ltx23_config = LTX23PipelineConfig()
         sana_wm_deployment = SanaWMPipelineConfig().get_model_deployment_config()
 
         self.assertIsNone(qwen_deployment.fsdp_auto_min_available_memory_gb)
@@ -816,6 +820,18 @@ class TestOffloadDefaults(unittest.TestCase):
         self.assertEqual(
             ltx_deployment.auto_disable_component_offload_components, ("dit",)
         )
+        self.assertEqual(
+            ltx_deployment.auto_cfg_parallel_degree_by_num_gpus, ((4, 1), (8, 1))
+        )
+        self.assertEqual(ltx_deployment.get_auto_cfg_parallel_degree(4), 1)
+        self.assertEqual(ltx_deployment.get_auto_cfg_parallel_degree(8), 1)
+        self.assertEqual(ltx_deployment.get_auto_cfg_parallel_degree(2), 2)
+        self.assertFalse(
+            LTX2PipelineConfig().dit_config.arch_config.enable_packed_qkv_input_a2a
+        )
+        self.assertFalse(
+            ltx23_config.dit_config.arch_config.enable_packed_qkv_input_a2a
+        )
 
         self.assertEqual(sana_wm_deployment.fsdp_auto_min_available_memory_gb, 60)
         self.assertTrue(sana_wm_deployment.auto_dit_layerwise_offload)
@@ -872,6 +888,24 @@ class TestOffloadDefaults(unittest.TestCase):
         self.assertFalse(args.use_fsdp_inference)
         self.assertFalse(args.enable_cfg_parallel)
 
+    def test_auto_ltx23_large_gpu_counts_prefer_sp_over_cfg_parallel(self):
+        for num_gpus in (4, 8):
+            with self.subTest(num_gpus=num_gpus):
+                args = self._from_dict_with_pipeline_config(
+                    LTX2PipelineConfig(),
+                    kwargs={
+                        "model_path": "Lightricks/LTX-2.3",
+                        "num_gpus": num_gpus,
+                        "performance_mode": "auto",
+                    },
+                )
+
+                self.assertFalse(args.enable_cfg_parallel)
+                self.assertEqual(args.cfg_parallel_degree, 1)
+                self.assertEqual(args.sp_degree, num_gpus)
+                self.assertEqual(args.ulysses_degree, num_gpus)
+                self.assertEqual(args.ring_degree, 1)
+
     def test_manual_mode_preserves_unset_performance_args(self):
         args = self._from_dict_with_pipeline_config(
             QwenImagePipelineConfig(),