minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mick
2026-06-27 19:23:14 +08:00
committed by GitHub
parent 91f9e7372d
commit 3306233961
11 changed files with 153 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/models/dits/ltx_2.py
+1
View File
@@ -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
+1
View File
@@ -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
+8
View File
@@ -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
+17 -3
View File
@@ -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:
python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py
+9
View File
@@ -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
+29 -7
View File
@@ -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,
python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py
+6
View File
@@ -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
+7 -4
View File
@@ -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(
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py
+7
View File
@@ -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
+33 -19
View File
@@ -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:
python/sglang/multimodal_gen/test/unit/test_server_args.py
+35 -1
View File
@@ -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(),