From 33062339616b1e09b1cf81e015575799dd4246c2 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 27 Jun 2026 19:23:14 +0800 Subject: [PATCH] [diffusion] optimize: optimize LTX2.3 CFG path (#28624) --- .../configs/models/dits/ltx_2.py | 1 + .../configs/pipeline_configs/ltx_2.py | 1 + .../model_deployment_config.py | 8 +++ .../runtime/layers/attention/layer.py | 20 +++++-- .../runtime/models/dits/ltx_2.py | 9 ++++ .../executors/parallel_executor.py | 36 ++++++++++--- .../runtime/pipelines_core/stages/base.py | 6 +++ .../ltx_2/decoding_av.py | 11 ++-- .../model_specific_stages/ltx_2/denoising.py | 7 +++ .../multimodal_gen/runtime/server_args.py | 52 ++++++++++++------- .../test/unit/test_server_args.py | 36 ++++++++++++- 11 files changed, 153 insertions(+), 34 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/models/dits/ltx_2.py b/python/sglang/multimodal_gen/configs/models/dits/ltx_2.py index 22b036f0c..15d59f5d0 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/ltx_2.py +++ b/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__() diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py index fb03fcd29..9c6e9b47e 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py +++ b/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): diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py b/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py index e01994075..378bb0529 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 3c8c8be95..0039b4935 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/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,9 +783,21 @@ class USPAttention(nn.Module): # Ulysses-style All-to-All for sequence/head sharding if sp_size > 1: # -> [B, S, H_local, D] - 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) + 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) # Ring Attention within subgroups or local attention if get_ring_parallel_world_size() > 1: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index fa1fbf2ac..e85c2da44 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py index 48ffee232..ebc5d14f1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/executors/parallel_executor.py @@ -68,16 +68,38 @@ class ParallelExecutor(PipelineExecutor): torch.distributed.barrier() elif paradigm == StageParallelismType.CFG_PARALLEL: - obj_list = [batch] if rank == 0 else [] - # `dist.broadcast(src=...)` expects a global rank for process groups. - broadcasted_list = broadcast_pyobj( - obj_list, - rank=get_world_rank(), - dist_group=cfg_group.cpu_group, - src=cfg_group.ranks[0], + 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, + rank=get_world_rank(), + 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, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py index 66076500e..5eb76fa9f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py +++ b/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. diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py index 9138d709b..8a7d6e866 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py +++ b/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( - device_type=current_platform.device_type, - dtype=audio_vae_dtype, - enabled=audio_vae_autocast_enabled, + 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( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py index 577b86aa2..9c8453cb9 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index aa3acab8e..4fc29be1c 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -821,26 +821,40 @@ 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 - if ( - self.performance_mode != "manual" - and deployment_config.auto_enable_cfg_parallel - and self.num_gpus >= 2 - and self.num_gpus % cfg_group_size == 0 - and sp_unspecified - and ulysses_unspecified - and ring_unspecified - and self._model_default_uses_cfg() - ): - self.enable_cfg_parallel = True - logger.info( - "Automatically enabled CFG parallel for %d GPUs. " - "Use --sp-degree / --ulysses-degree to use sequence " - "parallelism instead.", - self.num_gpus, - ) - else: + 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 + and self.num_gpus >= 2 + and self.num_gpus % cfg_group_size == 0 + and sp_unspecified + and ulysses_unspecified + and ring_unspecified + and self._model_default_uses_cfg() + ): + 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 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: + 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: diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index 9f0705f95..d839bc141 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/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(),